ABC097 C問題（K-th Substring）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）097 C問題をC++で解いてみました。ABC097は、2018年5月12日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

C問題 K-th Substring（Difficulty : 1108）

問題の詳細は、リンク先をご覧ください。

ABC097 C問題 K-th Substring

素直に解くと時間が足りません。AtCoder Problems による Difficulty は 1108 でした。

解答案

素直にすべての部分文字列を set コンテナ sub_s に格納します。最後に辞書順に $K$ 番目に小さい文字列を出力します。

#include <bits/stdc++.h>
using namespace std;

int main()
{
	string s;
	cin >> s;
	int k;
	cin >> k;

	set<string> sub_s;
	for (int i = 0; i < (int)s.length(); ++i) {
		for (int j = 1; i + j <= (int)s.length(); ++j) {
			sub_s.insert(s.substr(i, j));
		}
	}

	auto result = sub_s.begin();
	for (int i = 0; i < k - 1; ++i) {
		++result;
	}

	cout << *result << endl;

	return 0;
}

$|s| \leqq 50$ を満たすデータセットに正解するため、200点の部分点がもらえますが、8個のテストケースで TLE（Time Limit Exceeded）判定となり満点である300点はもらえません。

$N= |s|$ とします。部分文字列の要素は $N^2$ 個存在し、それらを set コンテナに挿入するために $O(N \log N)$ の計算量が必要となります（$O(\log N)$ ではないことに注意してください）。合わせて計算量は $O(N^3 \log N)$ となり、間に合いません。

$1 \leqq K \leqq 5$ という制約に注目します。制約が小さいため、$K$ 文字を超える部分文字列は辞書順で $K$ 番目より大きくなるため、set コンテナに格納する必要がありません。この場合、部分文字列の要素が $N K$ 個で、set コンテナに挿入するために $O(K \log N K)$ の計算量が必要となります。合わせて $O(N K^2 \log N K \doteqdot 2 \times 10^6)$ となり間に合います。

C++ プログラム例（ABC097C）

部分文字列の長さを最大5文字に制限します。前のプログラムとの差は13行目だけです。

以下が、C++プログラムです。

#include <bits/stdc++.h>
using namespace std;

int main()
{
	string s;
	cin >> s;
	int k;
	cin >> k;

	set<string> sub_s;
	for (int i = 0; i < (int)s.length(); ++i) {
		for (int j = 1; j <= k; ++j) {
			sub_s.insert(s.substr(i, j));
		}
	}

	auto result = sub_s.begin();
	for (int i = 0; i < k - 1; ++i) {
		++result;
	}

	cout << *result << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

最後に

問題の制約に少し違和感を覚えましたが、計算量に気付かず TLE 判定を受けました。解説をよみ、内容を理解することができました。勉強になりました。

引き続き ABC の問題を紹介していきます。