ABC215 C問題 One More aab aba baa を itertools使って解く - 競技プログラミング悪戦苦闘日記

ABC215 C問題を解説します。

問題

文字列 $S$ の各文字を並べ替えて作ることが可能な文字列を辞書順にすべて列挙したとき、前から $K$ 番目にくる文字列を求めてください。

難しいポイント

文字列の長さが最大で8なので、全パターン列挙しても $8! = 40320$ です。
あまり大きい数ではないので、列挙した後に並べ替えることもできます。
どうやって全パターン列挙するかがポイントなのかなと思います。

解き方の方針

文字列の長さが最大で8なので、全パターン列挙して、並べ替えて、K番目の文字列を取り出します。
全パターン列挙は公式解説のやり方もありますが、色々とあります。
私はpythonで順列列挙してくれるitertools.permutationsを使いました。

サンプルの"aab"の場合、順列列挙すると以下の6パターンが出てきます。

('a', 'a', 'b')
('a', 'b', 'a')
('a', 'a', 'b')
('a', 'b', 'a')
('b', 'a', 'a')
('b', 'a', 'a')

重複した文字があると、パターンにも重複するので、setを使って重複を取り除きます。

('b', 'a', 'a') 
('a', 'a', 'b') 
('a', 'b', 'a')

あとは全部文字列に変換してソートします。

['aab', 'aba', 'baa']

最後に配列の $K$ 番目を出力します。

実装

from itertools import permutations
S, K = input().split()
K = int(K)
S_arr = [s for s in S]
S_arr.sort()
patterns = permutations(S_arr, len(S_arr))
list_pat = list(set(patterns))
pat_str = []
for lp in list_pat:
    pat_str.append("".join(lp))
pat_str.sort()
print(pat_str[K - 1])