ARC_097_A / ABC_097_C - K-th Substring

[zeikar]

Problem link

https://atcoder.jp/contests/arc097/tasks/arc097_a

Problem Summary

문자열 s 의 모든 부분 문자열 중에서 k 번째로 작은 문자열을 구하는 문제.

Solution

일단 힌트는 k가 아주 작다는 것이다. (최대 5)

s 의 길이가 최대 5000이므로 s의 모든 부분 문자열을 구하는 것은 n^2이 되고 이론상 시간 안에 들어올 수 있다.
중복을 제거하기 위해 맵을 사용하고, 구한 모든 부분 문자열을 맵에 넣으면 쉽게 풀 수 있다.

처음 제출한 답은 그냥 단순히 모든 부분 문자열 중 k 번째를 구했는데 당연히 시간초과. 약간의 최적화를 해서 맵의 크기를 최대 k개 까지만 유지했더니 1초정도로 통과.

여기서 더 최적화를 할 수 있는데 k 번째의 부분 문자열의 길이는 최대 k여야만 한다는 것을 이용하면 된다.
길이가 k가 넘으면 그보다 작은 부분 문자열(prefix)보다 크게 되기 때문이다. (aaaaaa > aaaaa) 이다.

Source Code

#include <iostream>
#include <string>
#include <algorithm>
#include <map>
using namespace std;

int main() {
	string s;
	int k;

	cin >> s >> k;

	map<string, bool> substrs;
	for (int i = 0; i < s.length(); i++)
	{
		for (int j = i + 1; j <= min(i + k, (int)s.length()); j++)
		{
			substrs[s.substr(i, j - i)] = true;

			if (substrs.size() > k)
			{
				substrs.erase(--substrs.end());
			}
		}
	}

	for (auto const& str : substrs)
	{
		k--;

		if (k == 0)
		{
			cout << str.first << endl;
			break;
		}
	}
}