要素の追加(add)(AtCoder(競プロ)1日1問【7日目】【Python】)

AtCoder Beginner Contest 097 C問題

abc097.contest.atcoder.jp

競技プログラミングで使える関数を実際の問題を通して紹介いたします。

今回は「add（）」です。

add()はset関数で作成したものに対して要素を加えることができます。

今回の問題としては、入力した文字列の部分を切り取ってできる文字列を作って、
たくさんできた文字列の中でK番目に小さいものを出力するというもの。したがって、以下のように実装します。

s = input()
K = int(input())

substr = set()
for i in range(len(s)):
    for j in range(i+1, i+K+1):
        substr.add(s[i:j])

print(sorted(substr)[K-1])

コードの改善

競技プログラミングではプログラムの処理完了までの速度も重要な要素の一つですので、より短く簡潔なコードを書くことが重要です。

s = input()
K = int(input())

kouho = [s]# for x in range(K)]

for atama in range(len(s)):
    for offset in range(K):
        tes = s[atama:atama+offset+1]
        if tes in kouho:
            continue
        kouho.append(tes)
        kouho = sorted(kouho)
        if len(kouho) > K:
            kouho.pop()

print(kouho[K-1])