ABC220-E Distance on Large Perfect Binary Tree

問題へのリンク

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 2");
            //14
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("14142 17320");
            //11284501
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    const long Hou = 998244353;

    // 左部分木と右部分木の、高さ合計 = D での、左右の部分木の組み合わせ情報
    struct SubTreeInfoDef
    {
        internal long LeftHeight;
        internal long RightHeight;
        internal long ProdPatternCnt;
    }

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

        long N = wkArr[0];
        long D = wkArr[1];

        // ノード数[高さ]なDict
        var TreeCntDict = new Dictionary<long, long>();
        long CurrCnt = 1;
        for (long I = 0; I <= N - 1; I++) {
            TreeCntDict[I] = CurrCnt;
            CurrCnt *= 2;
            CurrCnt %= Hou;
        }

        // 高さ合計 = D での、左右の部分木の組み合わせ情報を列挙
        var SubTreeInfoList = new List<SubTreeInfoDef>();

        // 左部分木の高さでループ
        for (long I = 0; I <= D; I++) {
            long LeftHeight = I;
            long RightHeight = D - I;

            // 左部分木の高さ < 右部分木の高さとする
            if (RightHeight < LeftHeight) break;

            // 完全二分木の高さを超える場合は、不適
            if (TreeCntDict.ContainsKey(LeftHeight) == false) {
                continue;
            }
            if (TreeCntDict.ContainsKey(RightHeight) == false) {
                continue;
            }
            long LeftPatternCnt = TreeCntDict[LeftHeight];
            long RightPatternCnt = TreeCntDict[RightHeight];

            // ノード数が1より大きければ、2で割る
            if (LeftPatternCnt > 1) LeftPatternCnt *= DeriveGyakugen(2);
            if (RightPatternCnt > 1) RightPatternCnt *= DeriveGyakugen(2);
            LeftPatternCnt %= Hou;
            RightPatternCnt %= Hou;

            // 積の法則
            long ProdPatternCnt = LeftPatternCnt * RightPatternCnt;
            ProdPatternCnt %= Hou;

            // 左部分木の高さ < 右部分木の高さとしたので、2倍して調整
            if (RightHeight > LeftHeight) {
                ProdPatternCnt *= 2;
                ProdPatternCnt %= Hou;
            }

            SubTreeInfoDef WillAdd;
            WillAdd.LeftHeight = LeftHeight;
            WillAdd.RightHeight = RightHeight;
            WillAdd.ProdPatternCnt = ProdPatternCnt;
            SubTreeInfoList.Add(WillAdd);
        }

        // Pattern合計[左右の部分木の高さの最大値]な配列
        long[] CntArr = new long[N];
        foreach (SubTreeInfoDef EachSubTreeInfo in SubTreeInfoList) {
            long Ind = Math.Max(EachSubTreeInfo.LeftHeight, EachSubTreeInfo.RightHeight);
            if (Ind <= CntArr.GetUpperBound(0)) {
                CntArr[Ind] += EachSubTreeInfo.ProdPatternCnt;
                CntArr[Ind] %= Hou;
            }
        }
        long[] RunSum = (long[])CntArr.Clone();
        for (long I = 1; I <= RunSum.GetUpperBound(0); I++) {
            RunSum[I] += RunSum[I - 1];
            RunSum[I] %= Hou;
        }

        // LCAの高さを全探索
        long Answer = 0;
        for (long I = 0; I <= N - 1; I++) {
            long Height = N - 1 - I;
            if (RunSum[Height] == 0) break;

            long NodeCnt = TreeCntDict[I];
            Answer += NodeCnt * RunSum[Height];
            Answer %= Hou;
        }

        // 順列の左右入れ替えがあるので2倍する
        Answer *= 2;
        Answer %= Hou;
        Console.WriteLine(Answer);
    }

    // 引数の逆元を求める
    static Dictionary<long, long> mMemoGyakugenDict = new Dictionary<long, long>();
    static long DeriveGyakugen(long pLong)
    {
        if (mMemoGyakugenDict.ContainsKey(pLong)) {
            return mMemoGyakugenDict[pLong];
        }
        return mMemoGyakugenDict[pLong] = DeriveBekijyou(pLong, Hou - 2, Hou);
    }

    // 繰り返し2乗法で、(NのP乗) Mod Mを求める
    static long DeriveBekijyou(long pN, long pP, long pM)
    {
        long CurrJyousuu = pN % pM;
        long CurrShisuu = 1;
        long WillReturn = 1;

        while (true) {
            // 対象ビットが立っている場合
            if ((pP & CurrShisuu) > 0) {
                WillReturn = (WillReturn * CurrJyousuu) % pM;
            }

            CurrShisuu *= 2;
            if (CurrShisuu > pP) return WillReturn;
            CurrJyousuu = (CurrJyousuu * CurrJyousuu) % pM;
        }
    }
}

解説

最初に、完全二分木の根を高さ0として、
高さごとのノード数を求めてます。

次に、左右の部分木の高さの組み合わせごとの
場合の数を求めてます。

その後、LCAの高さを全探索して、解を計上してます。