A13 Close Pairs

問題へのリンク

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("7 10");
            WillReturn.Add("11 12 16 22 27 28 31");
            //11
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

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

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

        var InsWaveletTree = new WaveletTree(AArr);

        long Answer = 0;
        for (int I = 0; I <= UB - 1; I++) {
            int RangeSta = I + 1;
            int RangeEnd = UB;

            Answer += InsWaveletTree.RangeFreq(RangeSta, RangeEnd, AArr[I] + 1, AArr[I] + K);
        }
        Console.WriteLine(Answer);
    }
}

// ウェーブレット木クラス
internal class WaveletTree
{
    // 座圧後のInd[元々の値]なDict
    private Dictionary<int, int> mZaatuDict = new Dictionary<int, int>();

    // 元々の値[座圧後のInd]なArr
    private int[] mDistinctArr;

    // 完全二分木のノードのDict
    private Dictionary<int, WaveletTreeNodeDef> mTreeNodeDict =
        new Dictionary<int, WaveletTreeNodeDef>();

    // コンストラクタ
    internal WaveletTree(IEnumerable<int> pInitEnum)
    {
        // 座圧して、座圧後のInd[元々の値]なDictを設定
        var DistinctSet = new HashSet<int>(pInitEnum);
        mDistinctArr = DistinctSet.OrderBy(pX => pX).ToArray();
        for (int I = 0; I <= mDistinctArr.GetUpperBound(0); I++) {
            mZaatuDict[mDistinctArr[I]] = I;
        }

        // 座圧済の入力List
        var FirstValList = pInitEnum.Select(pX => mZaatuDict[pX]).ToList();

        // 深さ優先探索で完全二分木を作成
        var Stk = new Stack<WaveletTreeNodeDef>();

        var WillPush = new WaveletTreeNodeDef();
        WillPush.mNodeID = 0;
        WillPush.mLevel = 1;

        string BinStr = Convert.ToString(DistinctSet.Count - 1, 2);
        int CurrBit = (1 << (BinStr.Length - 1));

        WillPush.mCurrBit = CurrBit;
        WillPush.mValRangeMin = 0;
        WillPush.mValRangeMax = CurrBit * 2 - 1;
        WillPush.mValList = FirstValList;
        WillPush.MemberSet();

        Stk.Push(WillPush);
        while (Stk.Count > 0) {
            WaveletTreeNodeDef Popped = Stk.Pop();

            mTreeNodeDict[Popped.mNodeID] = Popped;

            if (Popped.mCurrBit <= 1) continue;

            var ValListLeft = new List<int>();
            var ValListRight = new List<int>();

            int MidVal = (Popped.mValRangeMin + Popped.mValRangeMax) / 2;
            // 左の部分木
            var WillPushLeft = new WaveletTreeNodeDef();
            int ChildNodeLeft = DeriveChildNode(Popped.mNodeID);
            WillPushLeft.mNodeID = ChildNodeLeft;
            WillPushLeft.mLevel = Popped.mLevel + 1;
            WillPushLeft.mCurrBit = Popped.mCurrBit / 2;
            WillPushLeft.mValRangeMin = Popped.mValRangeMin;
            WillPushLeft.mValRangeMax = MidVal;

            // 右の部分木
            var WillPushRight = new WaveletTreeNodeDef();
            WillPushRight.mNodeID = ChildNodeLeft + 1;
            WillPushRight.mLevel = Popped.mLevel + 1;
            WillPushRight.mCurrBit = Popped.mCurrBit / 2;
            WillPushRight.mValRangeMin = MidVal + 1;
            WillPushRight.mValRangeMax = Popped.mValRangeMax;

            foreach (int EachVal in Popped.mValList) {
                if (WillPushLeft.mValRangeMin <= EachVal && EachVal <= WillPushLeft.mValRangeMax) {
                    ValListLeft.Add(EachVal);
                }
                else {
                    ValListRight.Add(EachVal);
                }
            }
            WillPushLeft.mValList = ValListLeft;
            WillPushRight.mValList = ValListRight;
            WillPushLeft.MemberSet();
            WillPushRight.MemberSet();

            Stk.Push(WillPushLeft);
            Stk.Push(WillPushRight);
        }
        //DebugPrint();
    }

    // [0,End]までのXの出現回数を返す
    internal int Rank(int pEnd, int pX)
    {
        if (mZaatuDict.ContainsKey(pX) == false) return 0;

        // 座圧結果を取得
        int TargetVal = mZaatuDict[pX];

        // 根から完全二分木を探索
        WaveletTreeNodeDef CurrNode = mTreeNodeDict[0];

        int RangeEnd = pEnd;
        while (true) {
            int CurrBit = CurrNode.mCurrBit;

            int ChildNode = DeriveChildNode(CurrNode.mNodeID);
            if (Math.Sign(CurrBit & TargetVal) == 1) {
                ChildNode++;
            }

            // 葉ノードの場合
            if (mTreeNodeDict.ContainsKey(ChildNode) == false) {
                if (TargetVal % 2 == 0) {
                    return CurrNode.Rank(RangeEnd, 0);
                }
                else {
                    return CurrNode.Rank(RangeEnd, 1);
                }
            }

            if (Math.Sign(CurrBit & TargetVal) == 0) {
                RangeEnd = CurrNode.Rank(RangeEnd, 0) - 1;
            }
            else {
                RangeEnd = CurrNode.Rank(RangeEnd, 1) - 1;
            }

            CurrNode = mTreeNodeDict[ChildNode];
        }
    }

    // [Sta,End]までのXの出現回数を返す
    internal int Rank(int pSta, int pEnd, int pX)
    {
        if (pSta == 0) return Rank(pEnd, pX);

        int Cnt1 = Rank(pSta - 1, pX);
        int Cnt2 = Rank(pEnd, pX);
        return Cnt2 - Cnt1;
    }

    // [pSta,End]のK番目に大きい値を返す
    internal int kth_Largest(int pSta, int pEnd, int pK)
    {
        // 根から完全二分木を探索
        WaveletTreeNodeDef CurrNode = mTreeNodeDict[0];

        while (true) {
            int ChildNode = DeriveChildNode(CurrNode.mNodeID);

            // 現在区間の1の数
            int RangeOneCnt = CurrNode.Rank(pSta, pEnd, 1);

            // 葉ノードの場合
            if (mTreeNodeDict.ContainsKey(ChildNode) == false) {
                if (RangeOneCnt < pK) {
                    return mDistinctArr[CurrNode.mValRangeMin];
                }
                else {
                    return mDistinctArr[CurrNode.mValRangeMax];
                }
            }

            if (RangeOneCnt < pK) {
                // 左部分木に遷移する場合

                // 現在区間の0の数
                int RangeZeroCnt = (pEnd - pSta + 1) - RangeOneCnt;

                int PrevZeroCnt = 0;
                if (pSta > 0) {
                    PrevZeroCnt = CurrNode.Rank(pSta - 1, 0);
                }

                pK -= RangeOneCnt;
                pSta = PrevZeroCnt;
                pEnd = PrevZeroCnt + RangeZeroCnt - 1;
            }
            else {
                // 右部分木に遷移する場合

                ChildNode++;

                int PrevOneCnt = 0;
                if (pSta > 0) {
                    PrevOneCnt = CurrNode.Rank(pSta - 1, 1);
                }
                pSta = PrevOneCnt;
                pEnd = PrevOneCnt + RangeOneCnt - 1;
            }

            CurrNode = mTreeNodeDict[ChildNode];
        }
    }

    // [pSta,End]のK番目に小さい値を返す
    internal int kth_Smallest(int pSta, int pEnd, int pK)
    {
        int RangeCnt = pEnd - pSta + 1;
        return kth_Largest(pSta, pEnd, RangeCnt - pK + 1);
    }

    // [pSta,End]のMin以上、Max以下の値の個数を返す
    internal int RangeFreq(int pSta, int pEnd, int pMin, int pMax)
    {
        int Cnt1 = Rank(pSta, pEnd, pMax);
        int Cnt2 = RangeFreq(pSta, pEnd, pMin);
        int Cnt3 = RangeFreq(pSta, pEnd, pMax);
        return Cnt2 - Cnt3 + Cnt1;
    }

    // [pSta,End]のMin以上の値の個数を返す
    internal int RangeFreq(int pSta, int pEnd, int pMin)
    {
        int MinVal = ExecNibunhou_LowerBound(pMin, mDistinctArr);

        if (MinVal == -1) return 0;

        int WillReturn = 0;

        // 根から完全二分木を探索
        WaveletTreeNodeDef CurrNode = mTreeNodeDict[0];

        while (true) {
            int ChildNode = DeriveChildNode(CurrNode.mNodeID);

            // 現在区間の1の数
            int RangeOneCnt = CurrNode.Rank(pSta, pEnd, 1);

            // 現在区間の0の数
            int RangeZeroCnt = (pEnd - pSta + 1) - RangeOneCnt;

            // 葉ノードの場合
            if (mTreeNodeDict.ContainsKey(ChildNode) == false) {
                if (MinVal <= CurrNode.mValRangeMin) {
                    WillReturn += RangeZeroCnt;
                }
                if (MinVal <= CurrNode.mValRangeMax) {
                    WillReturn += RangeOneCnt;
                }
                return WillReturn;
            }

            if (mTreeNodeDict[ChildNode + 1].mValRangeMin >= MinVal) {
                // 左部分木に遷移する場合

                // 右部分木は、解に計上
                WillReturn += RangeOneCnt;

                int PrevZeroCnt = 0;
                if (pSta > 0) {
                    PrevZeroCnt = CurrNode.Rank(pSta - 1, 0);
                }

                pSta = PrevZeroCnt;
                pEnd = PrevZeroCnt + RangeZeroCnt - 1;
            }
            else {
                // 右部分木に遷移する場合

                ChildNode++;

                int PrevOneCnt = 0;
                if (pSta > 0) {
                    PrevOneCnt = CurrNode.Rank(pSta - 1, 1);
                }
                pSta = PrevOneCnt;
                pEnd = PrevOneCnt + RangeOneCnt - 1;
            }
            CurrNode = mTreeNodeDict[ChildNode];
        }
    }

    // 二分法で、Val以上で最小の値を持つ、添字を返す
    static int ExecNibunhou_LowerBound(int pVal, int[] pArr)
    {
        // 最後の要素がVal未満の特殊ケース
        if (pVal > pArr.Last()) {
            return -1;
        }
        // 最初の要素がVal以上の特殊ケース
        if (pVal <= pArr[0]) {
            return 0;
        }

        int L = 0;
        int R = pArr.GetUpperBound(0);

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

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

    // [pSta,End]でX未満の最大値(最大下界)を返す
    internal bool PrevValue(int pSta, int pEnd, int pX, out int pPrevValue)
    {
        pPrevValue = int.MinValue;
        int Cnt = RangeFreq(pSta, pEnd, pX);
        int K = 1 + Cnt;
        if (pEnd - pSta + 1 >= K) {
            pPrevValue = kth_Largest(pSta, pEnd, K);
            return true;
        }
        return false;
    }

    // [pSta,End]でX超えの最小値(最小上界)を返す
    internal bool NextValue(int pSta, int pEnd, int pX, out int pNextValue)
    {
        pNextValue = int.MinValue;
        int Cnt1 = pEnd - pSta + 1;
        int Cnt2 = RangeFreq(pSta, pEnd, pX);
        int Cnt3 = Rank(pSta, pEnd, pX);
        int MoreCnt = Cnt2 - Cnt3;
        int K = Cnt1 - MoreCnt + 1;

        if (pEnd - pSta + 1 >= K) {
            pNextValue = kth_Smallest(pSta, pEnd, K);
            return true;
        }
        return false;
    }

    // 親ノードの添字を取得
    private int DeriveParentNode(int pTarget)
    {
        return (pTarget - 1) / 2;
    }

    // 子ノードの添字(小さいほう)を取得
    private int DeriveChildNode(int pTarget)
    {
        return pTarget * 2 + 1;
    }

    // 完全二分木のデバッグ表示
    internal void DebugPrint()
    {
        //Func<string, IEnumerable<int>, string> IntEnumJoin = (pSeparater, pEnum) =>
        //{
        //    string[] StrArr = Array.ConvertAll(pEnum.ToArray(), pX => pX.ToString());
        //    return string.Join(pSeparater, StrArr);
        //};

        //foreach (var EachPair in mTreeNodeDict.OrderBy(pX => pX.Key)) {
        //    WaveletTreeNodeDef wkP = EachPair.Value;
        //    Console.WriteLine("■■■■■■■■■■");
        //    Console.WriteLine("ノードID={0},Level={1}", wkP.mNodeID, wkP.mLevel);
        //    Console.WriteLine("CurrBit={0},ValRangeMin={1},ValRangeMax={2}",
        //        wkP.mCurrBit, wkP.mValRangeMin, wkP.mValRangeMax);
        //    Console.WriteLine("mValArr={0}", IntEnumJoin(",", wkP.mValList));
        //    Console.WriteLine("mBitArr={0}", IntEnumJoin(",", wkP.mBitArr));
        //    Console.WriteLine("mZeroIndList={0}", IntEnumJoin(",", wkP.mZeroIndList));
        //    Console.WriteLine("mOneIndList={0}", IntEnumJoin(",", wkP.mOneIndList));
        //}
    }
}

// 完全二分木のノードごとの完備辞書のクラス
internal class WaveletTreeNodeDef
{
    internal int mNodeID; // ノードID (0始まり)
    internal int mLevel; // ノードのレベル
    internal int mCurrBit; // 見ているビット位置
    internal int mValRangeMin; // 対応するValの最小値
    internal int mValRangeMax; // 対応するValの最大値
    internal List<int> mValList; // 座圧した値のList
    internal int[] mBitArr; // ビット配列
    internal List<int> mZeroIndList; // ビット配列で0のIndのList
    internal List<int> mOneIndList;  // ビット配列で1のIndのList

    // mBitArr,mZeroIndList,mOneIndListの3つを設定する
    internal void MemberSet()
    {
        // 木の深さとValArrから、ビット配列を作る
        mBitArr = new int[mValList.Count];
        for (int I = 0; I <= mBitArr.GetUpperBound(0); I++) {
            mBitArr[I] = Math.Sign(mValList[I] & mCurrBit);
        }

        // IndのListを設定
        mZeroIndList = new List<int>();
        mOneIndList = new List<int>();
        for (int I = 0; I <= mBitArr.GetUpperBound(0); I++) {
            if (mBitArr[I] == 0) mZeroIndList.Add(I);
            if (mBitArr[I] == 1) mOneIndList.Add(I);
        }
    }

    // ValArrの[0,End]までのpBitの出現回数を返す
    internal int Rank(int pEnd, int pBit)
    {
        var IndListP = new List<int>();
        if (pBit == 0) IndListP = mZeroIndList;
        if (pBit == 1) IndListP = mOneIndList;

        int L = 0;
        int R = IndListP.Count + 1;

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

            bool IsOK = true;
            if (Mid > IndListP.Count) {
                IsOK = false;
            }
            else if (IndListP[Mid - 1] > pEnd) {
                IsOK = false;
            }

            if (IsOK) {
                L = Mid;
            }
            else {
                R = Mid;
            }
        }
        return L;
    }

    // ValArrの[Sta,End]までのpBitの出現回数を返す
    internal int Rank(int pSta, int pEnd, int pBit)
    {
        if (pSta == 0) return Rank(pEnd, pBit);

        int Cnt1 = Rank(pSta - 1, pBit);
        int Cnt2 = Rank(pEnd, pBit);
        return Cnt2 - Cnt1;
    }
}

解説

ウェーブレット木を使ってます。