Official

D - aab aba baa Editorial by en_translator


For the problem regarding with the lexicographical order, the formula is to decide the characters from the first to last in order. Now, in this problem, how can we determine whether the first character of the answer is a or b?:

This can be judged in the following way:

  • If there are more than or equal to \(K\) strings which starts with a, then it is a;
  • otherwise, it is b.

For each \(0 \leq i \leq A, 0 \leq j \leq B\), precalculate \(C(i, j)\), the total number of strings consisting with \(i\) number of a’s and \(j\) number of b’s. Since it is equal the number of ways to travel from \((0, 0)\) to \((i, j)\) by repeating advancing in the positive direction of \(x\) or \(y\) axis by \(1\), so it can be calculated with a simple DP (Dynamic Programming).

Let \(S(i, j, k)\) be “the \(k\)-th string in the lexicographical order among the strings consisting of \(i\) a’s and \(j\) b’s.” This can be found by splitting into the following cases:

  • If \(i = 0\), \(j\) repetition of b;
  • If \(j = 0\), \(i\) repetition of a;
  • If \(i > 0\) and \(j > 0\):
    • If \(C(i - 1, j) \geq k\), a followed by \(S(i - 1, j, k)\);
    • If \(C(i - 1, j) < k\), b followed by \(S(i, j - 1, k - C(i - 1, j))\).

This can be implemented (in C++) as a recursive function as follows:

#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;
}

Alternatively, it can be implemented (in C++) with non-recursive loops:

#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;
}

The total time complexity is \(O((A + B)^2)\) or \(O(AB)\).

posted:
last update: