K-th Substring [AtCoder Beginner Contest 097 C]

atcoder.jp

文字列 S の部分文字列のうち、K 番目に小さいものを出力せよ、という問題です。

解説にもあるように、部分点を取るだけなら非常に簡単です。
長さ1、長さ2、長さ3......長さ S.length() の部分文字列を全て求めてから、ソートして、K 番目に小さいものを出力すればよいです。

ただしこの解法は、S の長さが最大のとき、計算量が 約1200万+ソート となり、間に合いません。

制約に注目してみます。

・K は最大でも5
・S は異なる substring を K 個以上持つ

長さ5 までの部分文字列を列挙した時点で、少なくとも 長さ1～長さ5 の5個の部分文字列を得られます。これらの中には必ず小さいほうから K 番目の部分文字列が含まれています。よって、今回の問題では、長さ6以上の部分文字列を列挙する必要はありません。

S の長さが 5 未満の場合はどうでしょうか。
この場合でも、S の中に少なくとも K 個以上の部分文字列が含まれていることが保証されているので、
安心して 長さ1～長さS.length(<5) までの部分文字列を列挙してよいです。

つまりこの問題は、

・文字列 S の長さ min(K, S.length()) 以下の部分文字列のうち、K 番目に小さいものは何か

と言い換えることができます。

このときの計算量は 約10000+ソート となるので、余裕をもって間に合います。

void solve () {

	String s = next();
	int k = nextInt();

	TreeSet<String> set = new TreeSet<>();
	for (int i=1; i<=Math.min(s.length(), k); i++) {
		for (int j=0; j<s.length()-i+1; j++) {
			set.add(s.substring(j,j+i));
		}
	}

	ArrayList<String> list = new ArrayList<>(set);

	out.println(list.get(k-1));

}

Set を ArrayList に変換せずとも、拡張for文などで Set を回して K 番目の値を取り出すこともできます。