ABC202 D問題 aab aba baa をメモして解く

ABC202 D問題を解説します。
基本は公式解説と同じですが、具体例を書きながら解説します。

atcoder.jp

問題

A 個の a と B 個の b からなる長さ の文字列のうち、辞書順で K 番目のものを求めてください。

難しいポイント

AもBも最大で30なので、全量は列挙しきれません。
なので辞書順でKという情報から文字列を求めないといけませんが、どうやって。。。
何かしら閃くために具体的に書き出すとが見えては来ますが。。。

解き方の方針

まずは"a"か"b"を頭からつけて導くことを考えます。
具体的にA = 2、B = 3のケースを列挙してみます。
1. aabbb
2. ababb
3. abbab
4. abbba
5. baabb
6. babab
7. babba
8. bbaab
9. bbaba
10. bbbaa

書き出してみると、先頭にaがくるケースは、その後に続く文字のバターン数になっていることがわかります。
この場合では、残り4文字でのパターンの数です。つまり4つの中から1つのaが入る組み合わせの数です。
先頭にbがくるケースは、残り4文字でのパターン数。です。
なので先頭にaが来るのは、Kが以下の時になります。そうでない場合はbです。

1文字目が決まったので、2文字目の場合も同様に処理していきたいのですが、注意が必要です。
Kを具体的に示しながら、解説します。
K = 4だとしたら、1文字目は"a"です。2文字目を考える時に、さっきのように残り3文字でパターン数を考えます。
ここで注意が必要なのは、2文字目が"a"の時にそれ以降は"bbb"になるパターンです。
2文字目に"a"が来る場合は1パターン、"b"が来る場合はで3パターンです。

単純に数式では表現できないということがわかったので、2次元配列にメモします。
aの残りの数とbの残りの数が与えられた時のパターン数memo[a][b]を求めます。
aかbのどちらかが0の時は"aa"や"bbb"などの1パターンしかないので、1になります。
他の場合は、などの組み合わせの数になりますが、パスカルの三角形（二項定理）の要領で左隣と上の数字を足して求めます。
ja.wikipedia.org

もう1つ罠があるので先頭に"b"が来るパターンも考えます。
K=8のパターンをみていきます。
最初に"b"がくる場合は、2文字目を考えるときにKの値に注意が必要です。
K=8のまま考えると、a=2、b=2の時に8番目の数を選ぶということになってしまいます。
ここでは先程の先頭が"a"かどうかの閾値 4を考慮にいれて、という数にKを更新してあげます。

実装

たかだか30文字なので再帰処理で実装しました。

A, B, K = map(int, input().split())

# memo[a][b] aがa個, bがb個の時の場合の数
memo = [[0 for _ in range(31)] for _ in range(31)]
memo[0][0] = 1
for i in range(A + 1):
    for j in range(B + 1):
        if i == 0:
            memo[i][j] = 1
            continue
        if j == 0:
            memo[i][j] = 1
            continue
        memo[i][j] += memo[i][j - 1] + memo[i - 1][j]


# 残りa文字、残りb文字の時にk番目の文字列
def get_S(a, b, k):
    if a == 0:
        return "b" * b
    if b == 0:
        return "a" * a
    if memo[a - 1][b] >= k:
        return "a" + get_S(a - 1, b, k)
    if memo[a - 1][b] < k:
        return "b" + get_S(a, b - 1, k - memo[a - 1][b])


print(get_S(A, B, K))