ABC200-E Patisserie ABC 2

問題へのリンク

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 5");
            //1 2 2
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("1000000 1000000000000000000");
            //1000000 1000000 1000000
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("9 47");
            //3 1 4
        }
        else if (InputPattern == "Input4") {
            WillReturn.Add("1000000 219466283770414767");
            //469174 260221 367435
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long mN;
    static long mK;

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

        Console.WriteLine(Solve());
    }

    static long[] mDPResult2;

    static string Solve()
    {
        long UB = mN * 3;

        // 場合の数[和]なDP表
        long[] PrevDP = new long[UB + 1];
        PrevDP[0] = 1;

        for (long I = 1; I <= 3; I++) {
            long[] CurrDP = new long[UB + 1];
            for (long J = 0; J <= UB; J++) {
                if (PrevDP[J] == 0) continue;

                // いもす法で更新するIndのMinとMax
                long UpdateIndMin = J + 1;
                long UpdateIndMax = J + mN + 1;

                CurrDP[UpdateIndMin] += PrevDP[J];
                if (UpdateIndMax <= UB) {
                    CurrDP[UpdateIndMax] -= PrevDP[J];
                }
            }

            // 累積和に変更してから、すり替え
            for (long J = 1; J <= UB; J++) {
                CurrDP[J] += CurrDP[J - 1];
            }
            PrevDP = CurrDP;

            // DP結果を保存しておく
            if (I == 2) mDPResult2 = (long[])CurrDP.Clone();
        }

        long TargetSum = 0;
        long PatternCntSum1 = 0;
        long TargetRank = mK;

        // A+B+Cを決定する
        for (long I = 0; I <= UB; I++) {
            long PrevSum = PatternCntSum1;
            PatternCntSum1 += PrevDP[I];
            if (PatternCntSum1 >= TargetRank) {
                TargetSum = I;
                TargetRank -= PrevSum;
                break;
            }
        }

        // Aを決定する
        long DecideA = long.MaxValue;
        long PatternCntSum2 = 0;
        for (long LoopA = 1; LoopA <= TargetSum - 2; LoopA++) {
            // 下限のチェック
            if (TargetSum > LoopA + mN * 2) continue;

            long PrevSum = PatternCntSum2;
            PatternCntSum2 += DerivePatternCnt2(TargetSum, LoopA);
            if (PatternCntSum2 >= TargetRank) {
                DecideA = LoopA;
                TargetRank -= PrevSum;
                break;
            }
        }

        // Bを決定する
        long DecideB = long.MaxValue;
        for (long LoopB = 1; LoopB <= TargetSum - 2; LoopB++) {
            // 下限のチェック
            if (TargetSum > DecideA + LoopB + mN) continue;

            if (--TargetRank == 0) {
                DecideB = LoopB;
                break;
            }
        }

        long DecideC = TargetSum - DecideA - DecideB;
        return string.Format("{0} {1} {2}", DecideA, DecideB, DecideC);
    }

    // 和とAの値を引数にして、BとCの順列が何通りあるかを返す
    static long DerivePatternCnt2(long pSum, long pA)
    {
        long RestSum = pSum - pA;
        return mDPResult2[RestSum];
    }
}

解説

最初にDPで、
2個取った時の、場合の数[和]と
3個取った時の、場合の数[和]を求めておきます。

次に、
OrderBy A+B+C ASC , A ASC , B ASC
でソートした時のK番目が欲しいので
A+B+Cの昇順で件数の累積和を求め、累積和がK以上になったら、
その時のA+B+Cに注目します。
Kからは、1つ前の累積和を引いておきます。

次にAの昇順で件数の累積和を求め、累積和がK以上になったら、
その時のAに注目します。
Kからは、1つ前の累積和を引いておきます。

次にBとCを、線形オーダーでいいので、求めれば、解が分かります。