ACM-ICPC 2018 模擬国内予選: D. 短歌数

解法

「長さ$K$の数列であって数列$p$をprefixとして持ちかつ短歌数であるようなものの数$f(K, p)$」が求められれば先頭から順に決めていける(典型)。たぶん$O(\log N)$とかだと思う。

初手$10$進数展開して数列を考える。 とりあえず長さを固定したいので「長さ$K$の数列であって短歌数なものの数」を考え、では次に$\dots$とやれば終了。

二分探索という話も聞こえてきたが実装が面倒なだけでやることは同じな気がする。

実装

#!/usr/bin/env python3
def count_tanka_numbers_of_length(k):
    binomial_9_2 = 9 * 8 // 2
    pattern = 2 ** k - 2
    pattern_with_0 = 2 ** (k - 1) - 1
    return binomial_9_2 * pattern + 9 * pattern_with_0

def count_tanka_numbers_of_length_k_with_prefix(k, prefix):
    kind = len(set(prefix)) 
    if kind == 0:
        return count_tanka_numbers_of_length(k)
    elif kind == 1:
        if 0 in prefix:
            return 0
        pattern = 2 ** (k - len(prefix)) - 1
        return 9 * pattern
    elif kind == 2:
        return 2 ** (k - len(prefix))
    else:
        return 0

def solve(n):
    n -= 1
    acc = 0
    k = 1
    while not (n < acc + count_tanka_numbers_of_length(k)):
        acc += count_tanka_numbers_of_length(k)
        k += 1
    prefix = []
    for _ in range(k):
        for d in range(10):
            if not (n < acc + count_tanka_numbers_of_length_k_with_prefix(k, prefix + [ d ])):
                acc += count_tanka_numbers_of_length_k_with_prefix(k, prefix + [ d ])
            else:
                prefix += [ d ]
                break
        else:
            assert False
    return ''.join(map(str, prefix))

while True:
    n = int(input())
    if n == 0:
        break
    print(solve(n))