問題概要
N個の文字列が与えられ、これらの文字列から長さKの数列を作って文字列を連結します。全てのパターンで作られた文字列を辞書順に並べたとき、X番目の文字列を求める問題です。
具体的には、1からNまでの数字をK個選んで並べた数列(A₁, A₂, ..., Aₖ)に対して、文字列SA₁ + SA₂ + ... + SAₖを作ります。このような数列はN^K個存在し、それぞれから作られる文字列を辞書順にソートしてX番目を見つけます。
制約
- 1 ≤ N ≤ 10
- 1 ≤ K ≤ 5
- 1 ≤ X ≤ N^K
- 各文字列Siは英小文字からなる長さ10以下の文字列
入出力例
入力例1:
3 2 6
abc
xxx
abc
出力例1:
abcxxx
この例では、3つの文字列から長さ2の数列を作るので、3²=9通りの組み合わせができます。それぞれの組み合わせで作られる文字列を辞書順に並べると、6番目が「abcxxx」になります。
最初のアプローチと私の悩み
私が最初に書いたコードがこちら：
import itertools
n, k, x = map(int, input().split())
s_list = [input() for _ in range(n)]
index_list = [i for i in range(n)] * k
# print(index_list)
indexs_list = list(set(itertools.combinations(index_list, k)))
index_list.sort()
text_list = []
for indexs in indexs_list:
texts = []
for i in range(len(indexs)):
texts.append(s_list[indexs[i]])
texts = "".join(texts)
text_list.append(texts)
text_list.sort()
print(text_list[x - 1])
悩んだポイント
「全探索で間に合うのか？」を考えているうちに、何をしたいのかよくわからないコードになってしまいました。
制約を見ると最大10^5通りの組み合わせなので全探索で十分なはずなのに、なぜか複雑に考えすぎてしまったんです。特に、重複を許す順列を作りたかったのにcombinationsを使ってしまい、さらにsetで重複を除去するという謎の処理まで入れてしまいました。
私のコードのダメな点
元のコードには以下のような問題がありました：
まずitertools.combinationsを使っているのが根本的な間違いです。この問題では「重複を許す順列」が必要なのに、combinationsは重複を許さない組み合わせを生成します。
次にindex_list = [i for i in range(n)] * kで作ったリストをcombinationsに渡すという発想も複雑すぎました。もっとシンプルにproductを使えば一発で解決できます。
さらにlist(set(...))で重複除去をしていますが、これは完全に不要な処理です。
最後に変数名もindexs_list、textsなど、意図が伝わりにくいものになってしまいました。
改善されたコード
問題点を整理した上で、シンプルで分かりやすいコードに書き直しました：
from itertools import product
# 入力を読み取る
N, K, X = map(int, input().split())
strings = []
for i in range(N):
strings.append(input().strip())
# 全ての組み合わせを生成して対応する文字列を作成
all_strings = []
# 1からNまでの数字を使って長さKの重複順列を生成
for combination in product(range(1, N + 1), repeat=K):
# 各組み合わせに対応する文字列を連結
concatenated = ""
for index in combination:
# インデックスは1ベースなので、0ベースに変換
concatenated += strings[index - 1]
# 作成した文字列をリストに追加
all_strings.append(concatenated)
# 辞書順にソート
all_strings.sort()
# X番目の文字列を出力（1ベースなので-1）
print(all_strings[X - 1])
元コードとの違い
最も大きな違いはitertools.productを使ったことです。元のコードではcombinationsを使っていましたが、これは重複を許さない組み合わせしか作れません。productなら重複を許す順列を簡単に生成できます。
次に不要な処理を全て削除しました。元のコードにあったsetによる重複除去や、複雑なリスト生成処理は一切必要ありませんでした。
また、変数名も意図が分かりやすいものに変更しています。strings、all_strings、concatenatedなど、何を格納しているのかが一目で分かる名前にしました。
なぜこの方法が良いのか
まずitertools.productは重複順列を生成するための標準的な方法です。product(range(1, N+1), repeat=K)とするだけで、1からNまでの数字を使ったK個の重複順列を全て生成できます。
コードの流れもとてもシンプルになりました。「組み合わせを生成→文字列を連結→ソート→X番目を取得」という処理が直線的に書かれているので、読みやすく理解しやすいコードになっています。
さらに計算量的にも効率的です。余計な重複除去処理がないので、O(N^K)の時間で全ての組み合わせを生成できます。
1-3行目で入力を読み取ります
N, K, Xを一度に読み取り、N個の文字列をリストに格納しています。
6-15行目では全組み合わせを生成して文字列を作成します
product(range(1, N + 1), repeat=K)で1からNまでの数字を使った長さKの重複順列を生成します。各組み合わせは(1, 2)のような形で得られます。combinationの各要素indexに対して、strings[index - 1]で対応する文字列を取得しています。ここで1ベースから0ベースに変換しているわけですね。全ての文字列を連結してall_stringsリストに追加していきます。
18行目で辞書順にソートしています
Pythonのsort()メソッドは文字列を辞書順に自動的にソートしてくれます。
21行目で結果を出力します
X番目は1ベースなので、0ベースのリストではX-1番目を取得する必要があります。
計算量
- 時間計算量：O(N^K × L × log(N^K))
- N^K個の組み合わせを生成：O(N^K)
- 各文字列の長さをLとして連結：O(L)
- ソート：O(N^K × log(N^K))
- 空間計算量：O(N^K × L)
- N^K個の文字列を格納
制約上、N≤10、K≤5、文字列長≤10なので、最悪でも10⁵個の文字列をソートする程度で十分高速です。
学んだこと
「combinations vs product」の選択
最初はitertools.combinationsを使いたくなるかもしれませんが、これは重複を許さない組み合わせしか作れません。この問題では同じ文字列を何度でも使えるので、重複を許すitertools.productを選ぶ必要があります。
計算量への過度な心配
制約を見て「本当に全探索で大丈夫？」と心配になることがあります。でも冷静に計算してみると、最大でも10^5通りなので全く問題ありません。むしろ複雑に考えすぎると、分かりにくいコードになってしまいます。
インデックスの混乱
問題文では文字列に1, 2, 3...という番号が付いていますが、Pythonのリストは0から始まります。この変換でミスしやすいので注意してください。
変数名の付け方
デバッグ中に変数名が適当になりがちですが、後から見直したときに何をしているのか分からなくなってしまいます。最初から意図の分かる名前を付けておくことをお勧めします。
まとめ
この問題は制約が小さいため、全ての組み合わせを生成してソートするという直接的なアプローチが有効でした。より制約が大きい場合は、辞書順でX番目を直接計算するアルゴリズムが必要になりますが、AtCoderのC問題レベルでは全探索で解ける問題が多いです。