ABC429-D On AtCoder Conference

問題へのリンク

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("5 3 2");
            WillReturn.Add("1 2 1 0 1");
            //9
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("1 1000000000000 1");
            WillReturn.Add("1");
            //1000000000000
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long[] GetSplitArr(string pStr)
    {
        return (pStr == "" ? new string[0] : pStr.Split(' ')).Select(pX => long.Parse(pX)).ToArray();
    }

    static Fenwick_Tree mIns_Fenwick_Tree;
    static long mFenwick_Tree_UB;
    static long mStopCnt;

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] wkArr = GetSplitArr(InputList[0]);
        long M = wkArr[1];
        mStopCnt = wkArr[2];

        long[] AArr = GetSplitArr(InputList[1]);
        var CntDict = new Dictionary<long, long>();
        foreach (long EachA in AArr) {
            if (CntDict.ContainsKey(EachA) == false) {
                CntDict[EachA] = 0;
            }
            CntDict[EachA]++;
        }

        if (CntDict.ContainsKey(0) == false) {
            CntDict[0] = 0;
        }
        if (CntDict.ContainsKey(M - 1) == false) {
            CntDict[M - 1] = 0;
        }

        // 座標圧縮して、フェニック木を作成
        Dictionary<long, long> ZaatuDict = DeriveZaatuDict(CntDict.Keys);

        var NewCntDict = new Dictionary<long, long>();
        foreach (var EachPair in ZaatuDict) {
            NewCntDict[EachPair.Value] = CntDict[EachPair.Key];
        }

        mIns_Fenwick_Tree = new Fenwick_Tree(ZaatuDict.Count - 1);
        mFenwick_Tree_UB = mIns_Fenwick_Tree.GetUB();
        foreach (var EachPair in NewCntDict) {
            mIns_Fenwick_Tree[EachPair.Key] = EachPair.Value;
        }

        long[] KeysArr = CntDict.Keys.ToArray();
        Array.Sort(KeysArr);
        long Answer = 0;
        for (long I = 0; I <= KeysArr.GetUpperBound(0); I++) {
            long Diff;
            if (I < KeysArr.GetUpperBound(0)) {
                long NextInd = KeysArr[I + 1];
                Diff = NextInd - KeysArr[I];
            }
            else {
                Diff = 1;
            }

            long RangeSta = I + 1;
            if (RangeSta > KeysArr.GetUpperBound(0)) {
                RangeSta = 0;
            }

            // フェニック木から取得する区間長が1で解の場合
            long FirstRangeSum = GetRangeSum(RangeSta, RangeSta);
            if (FirstRangeSum >= mStopCnt) {
                Answer += FirstRangeSum * Diff;
                continue;
            }

            // フェニック木から取得する区間長を二分探索
            long L = 1;
            long R = NewCntDict.Count;

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

                if (GetRangeSum(RangeSta, RangeSta + Mid - 1) >= mStopCnt) {
                    R = Mid;
                }
                else {
                    L = Mid;
                }
            }
            long AnswerRangeSum = GetRangeSum(RangeSta, RangeSta + R - 1);
            Answer += AnswerRangeSum * Diff;
        }
        Console.WriteLine(Answer);
    }

    // フェニック木の区間和を返す
    static long GetRangeSum(long pRangeSta, long pRangeEnd)
    {
        var RangeInfoList = ArrUtil.GetRangeInfoList(pRangeSta, pRangeEnd, mFenwick_Tree_UB);
        long SumVal = 0;
        foreach (ArrUtil.RangeInfoDef EachRangeInfo in RangeInfoList) {
            SumVal += mIns_Fenwick_Tree.GetSum(EachRangeInfo.IndSta, EachRangeInfo.IndEnd);
        }
        return SumVal;
    }

    //////////////////////////////////////////////////////////////////////////
    // 列挙を引数として、座標圧縮し、座圧後の値[座圧前の値]なDictを返す
    //////////////////////////////////////////////////////////////////////////
    static Dictionary<long, long> DeriveZaatuDict(IEnumerable<long> pEnum)
    {
        var ZaatuDict = new Dictionary<long, long>();
        var ValSet = new HashSet<long>(pEnum);
        long No = 0;
        foreach (long EachVal in ValSet.OrderBy(pX => pX)) {
            ZaatuDict[EachVal] = No;
            No++;
        }
        return ZaatuDict;
    }
}

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

    // ノードのIndexの列挙を返す
    internal IEnumerable<long> GetNodeIndEnum()
    {
        for (long I = 0; I <= mExternalArrUB; I++) {
            yield return I;
        }
    }

    // 木のノードのUBを返す
    internal long GetUB()
    {
        return mExternalArrUB;
    }

    // コンストラクタ(外部配列のUBのみ指定)
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

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

    // コンストラクタ(初期化用の配列指定)
    internal Fenwick_Tree(long[] pArr)
        : this(pArr.GetUpperBound(0))
    {
        for (long I = 0; I <= pArr.GetUpperBound(0); I++) {
            this.Add(I, pArr[I]);
        }
    }

    // コンストラクタ(初期化用のList指定)
    internal Fenwick_Tree(List<long> pList)
        : this(pList.Count - 1)
    {
        for (int I = 0; I <= pList.Count - 1; I++) {
            this.Add(I, pList[I]);
        }
    }

    // インデクサ
    internal long this[long pInd]
    {
        get { return GetSum(pInd, pInd); }
        set { Add(pInd, value - GetSum(pInd, pInd)); }
    }

    // [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

#region ArrUtil
// 配列のユーティリティクラス
internal class ArrUtil
{
    // 区間情報
    internal struct RangeInfoDef
    {
        internal long IndSta;
        internal long IndEnd;
    }

    // IndStaとIndEndとUBを引数とし、
    // IndEndがUBを超えたら、0から開始に変更し、
    // 区間のListを返す
    static internal List<RangeInfoDef> GetRangeInfoList(long pIndSta, long pIndEnd, long UB)
    {
        var WillReturn = new List<RangeInfoDef>();
        if (pIndEnd <= UB) {
            RangeInfoDef WillAdd1;
            WillAdd1.IndSta = pIndSta;
            WillAdd1.IndEnd = pIndEnd;
            WillReturn.Add(WillAdd1);
        }
        else {
            // IndEndがUBを超えてる場合
            RangeInfoDef WillAdd1;
            WillAdd1.IndSta = pIndSta;
            WillAdd1.IndEnd = UB;
            WillReturn.Add(WillAdd1);

            long Rest = pIndEnd - UB;
            RangeInfoDef WillAdd2;
            WillAdd2.IndSta = 0;
            WillAdd2.IndEnd = Rest - 1;
            WillReturn.Add(WillAdd2);
        }
        return WillReturn;
    }
}
#endregion

解説

座標圧縮して、フェニック木を作成し、
始点ごとに、どこまでの区間で、C以上になるかを二分探索してます。

解への寄与は、座圧してない状態で前後の差分から求めることができます。