ABC342-D Square Pair

問題へのリンク

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");
            WillReturn.Add("0 3 2 8 12");
            //6
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("8");
            WillReturn.Add("2 2 4 6 3 100 100 25");
            //7
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static List<long> mHeihousuuList = new List<long>();

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

        DeriveHeihousuuList(AArr.Max());

        for (long I = 0; I <= UB; I++) {
            if (AArr[I] <= 1) continue;
            for (int J = 0; J <= mHeihousuuList.Count - 1; J++) {
                while (AArr[I] % mHeihousuuList[J] == 0) {
                    AArr[I] /= mHeihousuuList[J];
                }
            }
        }

        var InsWaveletTree = new WaveletTree(AArr);
        long Answer = 0;
        for (long I = 0; I <= UB - 1; I++) {
            if (AArr[I] == 0) {
                Answer += UB - I;
            }
            else {
                Answer += InsWaveletTree.Rank(I + 1, UB, 0);
                Answer += InsWaveletTree.Rank(I + 1, UB, AArr[I]);
            }
        }
        Console.WriteLine(Answer);
    }

    // Limit以下の平方数のListを求める
    static void DeriveHeihousuuList(long pLimit)
    {
        for (long I = 2; I * I <= pLimit; I++) {
            mHeihousuuList.Add(I * I);
        }
    }
}

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

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

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

    // コンストラクタ
    internal WaveletTree(IEnumerable<long> pInitEnum)
    {
        // 座圧して、座圧後のInd[元々の値]なDictを設定
        var DistinctSet = new HashSet<long>(pInitEnum);
        mDistinctArr = DistinctSet.OrderBy(pX => pX).ToArray();
        for (long 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);
        long 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<long>();
            var ValListRight = new List<long>();

            long MidVal = (Popped.mValRangeMin + Popped.mValRangeMax) / 2;
            // 左の部分木
            var WillPushLeft = new WaveletTreeNodeDef();
            long 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 (long 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 long Rank(long pEnd, long pX)
    {
        if (mZaatuDict.ContainsKey(pX) == false) return 0;

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

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

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

            long 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 long Rank(long pSta, long pEnd, long pX)
    {
        if (pSta == 0) return Rank(pEnd, pX);

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

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

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

            // 現在区間の1の数
            long 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の数
                long RangeZeroCnt = (pEnd - pSta + 1) - RangeOneCnt;

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

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

                ChildNode++;

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

            CurrNode = mTreeNodeDict[ChildNode];
        }
    }

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

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

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

        if (MinVal == -1) return 0;

        long WillReturn = 0;

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

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

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

            // 現在区間の0の数
            long 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;

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

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

                ChildNode++;

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

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

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

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

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

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

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

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

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

    // 子ノードの添字(小さいほう)を取得
    private long DeriveChildNode(long 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 long mNodeID; // ノードID (0始まり)
    internal long mLevel; // ノードのレベル
    internal long mCurrBit; // 見ているビット位置
    internal long mValRangeMin; // 対応するValの最小値
    internal long mValRangeMax; // 対応するValの最大値
    internal List<long> mValList; // 座圧した値のList
    internal long[] mBitArr; // ビット配列
    internal List<long> mZeroIndList; // ビット配列で0のIndのList
    internal List<long> mOneIndList;  // ビット配列で1のIndのList

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

        // IndのListを設定
        mZeroIndList = new List<long>();
        mOneIndList = new List<long>();
        for (long 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 long Rank(long pEnd, long pBit)
    {
        var IndListP = new List<long>();
        if (pBit == 0) IndListP = mZeroIndList;
        if (pBit == 1) IndListP = mOneIndList;

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

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

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

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

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

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

解説

2以上の平方数で割れるだけ割った配列で
ウェーブレット木を作り
●片方が0
●値が一致する
なペアの個数を求めてます。