ABC406-E Popcount Sum 3

問題へのリンク

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("2");
            WillReturn.Add("20 2");
            WillReturn.Add("1234567890 17");
            //100
            //382730918
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    const long Hou = 998244353;

    static void Main()
    {
        List<string> InputList = GetInputList();

        long[] wkArr = { };
        Action<string> SplitAct = pStr =>
            wkArr = pStr.Split(' ').Select(pX => long.Parse(pX)).ToArray();

        foreach (string EachStr in InputList.Skip(1)) {
            SplitAct(EachStr);
            long N = wkArr[0];
            long K = wkArr[1];
            long Result = Solve(N, K);
            Console.WriteLine(Result);
        }
    }

    static long Solve(long pN, long pK)
    {
        string BinStr = Convert.ToString(pN, 2);
        char[] BinArr = BinStr.ToCharArray();

        long PopCntUB = pK;

        // 場合の数[N未満確定フラグ,PopCnt]なDP表
        long[,] PrevDP1 = new long[2, PopCntUB + 1];
        PrevDP1[0, 0] = 1;

        // 合計[N未満確定フラグ,PopCnt]なDP表
        long[,] PrevDP2 = new long[2, PopCntUB + 1];
        PrevDP2[0, 0] = 0;

        long UB = BinArr.GetUpperBound(0);

        var CurrBitDict = new Dictionary<long, long>();
        long Curr2Beki = 1;
        for (long I = UB; 0 <= I; I--) {
            CurrBitDict[I] = Curr2Beki;
            Curr2Beki *= 2;
            Curr2Beki %= Hou;
        }

        for (long I = 0; I <= UB; I++) {
            long[,] CurrDP1 = new long[2, PopCntUB + 1];
            long[,] CurrDP2 = new long[2, PopCntUB + 1];
            for (long J = 0; J <= 1; J++) {
                for (long K = 0; K <= PopCntUB; K++) {
                    if (PrevDP1[J, K] == 0) continue;

                    for (char NewChar = '0'; NewChar <= '1'; NewChar++) {
                        if (J == 0 && BinArr[I] < NewChar) continue;

                        long NewJ = J;
                        if (BinArr[I] > NewChar) NewJ = 1;

                        int CurrNum = NewChar - '0';

                        long NewK = K;
                        if (NewChar == '1') {
                            NewK++;
                        }
                        if (NewK <= PopCntUB) {
                            CurrDP1[NewJ, NewK] += PrevDP1[J, K];
                            CurrDP1[NewJ, NewK] %= Hou;

                            if (NewChar == '1') {
                                CurrDP2[NewJ, NewK] += CurrBitDict[I] * PrevDP1[J, K];　// 寄与分を加算
                                CurrDP2[NewJ, NewK] %= Hou;
                            }
                            CurrDP2[NewJ, NewK] += PrevDP2[J, K]; // 元々の分を加算
                            CurrDP2[NewJ, NewK] %= Hou;
                        }
                    }
                }
            }
            PrevDP1 = CurrDP1;
            PrevDP2 = CurrDP2;
        }

        long Answer = 0;
        Answer += PrevDP2[0, PopCntUB];
        Answer %= Hou;
        Answer += PrevDP2[1, PopCntUB];
        Answer %= Hou;
        return Answer;
    }
}

解説

場合の数と合計
の二つを用意して桁DPしてます。

類題

ABC235-F Variety of Digits