AtCoderのABC    次のABCの問題へ    前のABCの問題へ

ABC347-E Set Add Query


問題へのリンク


C#のソース

using System;
using System.Collections.Generic;
using System.Linq;

class Program
{
    static string InputPattern = "InputX";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();

        if (InputPattern == "Input1") {
            WillReturn.Add("3 4");
            WillReturn.Add("1 3 3 2");
            //6 2 2
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("4 6");
            WillReturn.Add("1 2 3 2 4 2");
            //15 9 12 7
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] wkArr = InputList[0].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        long N = wkArr[0];

        long[] XArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        long UB = XArr.GetUpperBound(0);

        var Ins_Fenwick_Tree = new Fenwick_Tree(UB);

        // フェニック木に値を設定
        var XSet = new HashSet<long>();
        for (long I = 0; I <= UB; I++) {
            if (XSet.Contains(XArr[I])) {
                XSet.Remove(XArr[I]);
            }
            else {
                XSet.Add(XArr[I]);
            }
            Ins_Fenwick_Tree.Add(I, XSet.Count);
        }

        // IndList[値]なDict
        var IndListDict = new Dictionary<long, List<long>>();

        for (long I = 0; I <= UB; I++) {
            if (IndListDict.ContainsKey(XArr[I]) == false) {
                IndListDict[XArr[I]] = new List<long>();
            }
            IndListDict[XArr[I]].Add(I);
        }

        // 高速にシュミレーション
        long[] AnswerArr = new long[N + 1];

        XSet.Clear();
        for (long I = 0; I <= UB; I++) {
            if (XSet.Contains(XArr[I])) {
                XSet.Remove(XArr[I]);
                continue;
            }
            XSet.Add(XArr[I]);

            List<long> CurrList = IndListDict[XArr[I]];

            long ResultInd = ExecNibunhou_UpperBound(I, CurrList);
            long RangeSta = I;
            long RangeEnd = UB;

            if (ResultInd > -1) {
                RangeEnd = CurrList[(int)ResultInd] - 1;
            }
            AnswerArr[XArr[I]] += Ins_Fenwick_Tree.GetSum(RangeSta, RangeEnd);
        }

        var AnswerList = new List<long>();
        for (long I = 1; I <= N; I++) {
            AnswerList.Add(AnswerArr[I]);
        }

        Console.WriteLine(LongEnumJoin(" ", AnswerList));
    }

    // セパレータとLong型の列挙を引数として、結合したstringを返す
    static string LongEnumJoin(string pSeparater, IEnumerable<long> pEnum)
    {
        string[] StrArr = Array.ConvertAll(pEnum.ToArray(), pX => pX.ToString());
        return string.Join(pSeparater, StrArr);
    }

    // 二分法で、Val超えで最小の値を持つ、添字を返す
    static long ExecNibunhou_UpperBound(long pVal, List<long> pList)
    {
        // 最後の要素がVal以下の特殊ケース
        if (pVal >= pList.Last()) {
            return -1;
        }
        // 最初の要素がVal超えの特殊ケース
        if (pVal < pList[0]) {
            return 0;
        }

        long L = 0;
        long R = pList.Count - 1;

        while (L + 1 < R) {
            long Mid = (L + R) / 2;

            if (pList[(int)Mid] > pVal) {
                R = Mid;
            }
            else {
                L = Mid;
            }
        }
        return R;
    }
}

// フェニック木
#region Fenwick_Tree
internal class Fenwick_Tree
{
    private long[] mBitArr;
    private long mExternalArrUB;

    // コンストラクタ
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

        // フェニック木の外部配列は0オリジンで、
        // フェニック木の内部配列は1オリジンなため、2を足す
        mBitArr = new long[pExternalArrUB + 2];
    }

    // [pSta,pEnd] のSumを返す
    internal long GetSum(long pSta, long pEnd)
    {
        return GetSum(pEnd) - GetSum(pSta - 1);
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long pEnd)
    {
        pEnd++; // 1オリジンに変更

        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // [I] に Xを加算
    internal void Add(long pI, long pX)
    {
        pI++; // 1オリジンに変更

        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion


解説

考察すると、
集合の要素数を設定したフェニック木を用意し、
集合のAddされてからRemoveされるまでを二分探索で求め、
フェニック木の区間和を解に計上すれば良いと分かります。