Official

D - aab aba baa Editorial by KoD


辞書順が関係する問題では、前の文字から決めていくというのが定石です。では、この問題において、答えの先頭の文字が ab のどちらであるか判定するためにはどのようにすれば良いでしょうか?

これは次のように判定できます:

  • 先頭が a であるような文字列の総数が \(K\) 個以上なら a
  • そうでないなら b

\(0 \leq i \leq A, 0 \leq j \leq B\) について、a\(i\) 個と b\(j\) 個からなる文字列の総数 \(C(i, j)\) を前もって計算しておきます。これは \((0, 0)\) からスタートし、\(x\) 軸の正の方向あるいは \(y\) 軸の正の方向に \(1\) 進むことを繰り返して \((i, j)\) に至る方法の総数に等しいため、単純な動的計画法によって求めることが出来ます。

\(S(i, j, k)\) を「a\(i\) 個と b\(j\) 個からなる文字列のうち、辞書順で \(k\) 番目のもの」と定義します。これは次のような場合分けによって求められます:

  • \(i = 0\) のとき、b\(j\) 個からなる文字列
  • \(j = 0\) のとき、a\(i\) 個からなる文字列
  • \(i > 0\) かつ \(j > 0\) のとき
    • \(C(i - 1, j) \geq k\) のとき、a\(S(i - 1, j, k)\) を繋げた文字列
    • \(C(i - 1, j) < k\) のとき、b\(S(i, j - 1, k - C(i - 1, j))\) を繋げた文字列

再帰関数として実装すると次のようになります (C++):

#include <iostream>
#include <string>

constexpr int MAX = 30;
long long dp[MAX + 1][MAX + 1];

std::string find_kth(int A, int B, long long K) {
    if (A == 0) {
        return std::string(B, 'b');
    }
    if (B == 0) {
        return std::string(A, 'a');
    }
    if (K <= dp[A - 1][B]) {
        return std::string("a") + find_kth(A - 1, B, K);
    }
    else {
        return std::string("b") + find_kth(A, B - 1, K - dp[A - 1][B]);
    }
}

int main() {
    int A, B;
    long long K;
    std::cin >> A >> B >> K;
    dp[0][0] = 1;
    for (int i = 0; i <= A; ++i) {
        for (int j = 0; j <= B; ++j) {
            if (i > 0) {
                dp[i][j] += dp[i - 1][j];
            }
            if (j > 0) {
                dp[i][j] += dp[i][j - 1];
            }
        }
    }
    std::cout << find_kth(A, B, K) << '\n';
    return 0;
}

ループによって非再帰で実装することも可能です (C++):

#include <iostream>
#include <vector>
#include <string>

int main() {
    int A, B;
    long long K;
    std::cin >> A >> B >> K;
    std::vector<std::vector<long long>> dp(A + 1, std::vector<long long>(B + 1));
    dp[0][0] = 1;
    for (int i = 0; i <= A; ++i) {
        for (int j = 0; j <= B; ++j) {
            if (i > 0) {
                dp[i][j] += dp[i - 1][j];
            }
            if (j > 0) {
                dp[i][j] += dp[i][j - 1];
            }
        }
    }
    while (A > 0 and B > 0) {
        if (K <= dp[A - 1][B]) {
            std::cout << 'a';
            A -= 1;
        }
        else {
            K -= dp[A - 1][B];
            std::cout << 'b';
            B -= 1;
        }
    }
    std::cout << std::string(A, 'a');
    std::cout << std::string(B, 'b');
    std::cout << '\n';
    return 0;
}

全体の計算量は \(O((A + B)^2)\) あるいは \(O(AB)\) です。

posted:
last update: