ACM ICPC NEERC 2016: Problem L. List of Primes

neerc2016-L.cpp

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

pair<int, int> get_next_prime(int k)
{
	for (int i = 2; i < k; i++)
		if (k % i == 0)
			return get_next_prime(k + 1);

	int res = 0;
	for (int i = k; i > 0; i /= 10)
		res++;

	return make_pair(k, res);
}

struct A {
	long long count; 
	long long length;
};
map<int, map<int, A>> mydata;

void print_sub(const string& prefix, long long skip, long long take, int sum = 2, int min = 0)
{
	for (auto& m : mydata[sum]) {
		if (take == 0)
			return;

		if (m.first <= min)
			continue;

		long long len = m.second.length + m.second.count*(prefix.length() - 1);
		if (len < skip) {
			skip -= len;
			continue;
		}

		if (m.first == sum)
			cout << (prefix + to_string(sum) + "], ").substr(skip, take);
		else
			print_sub(prefix + to_string(m.first) + ", ", skip, take, sum - m.first, m.first);

		take = max(take - (len - skip), 0ll);
		skip = 0;
	}

	if (min == 0 && take > 0)
		print_sub(prefix, skip, take, sum + 1);
}

int main()
{
	long long n, k;
	cin >> n >> k;

	long long full_length = 0;
	vector<pair<int, int>> primes = { make_pair(2, 1) };
	for (int sum = 2; true; sum++) {
		if (primes.back().first < sum)
			primes.emplace_back(get_next_prime(sum));

		for (auto& p : primes)
		{
			int prime, len;
			tie(prime, len) = p;

			if (prime > sum)
				break;

			if (prime == sum)
				mydata[prime][prime] = A{ 1,  len + 4 };

			for (auto& v : mydata[sum - prime]) {
				if (v.first > prime) {
					auto& t = mydata[sum][prime];
					t.count += v.second.count;
					t.length += v.second.length + (2 + len)*v.second.count;
				}
			}
		}

		for (auto& v : mydata[sum])
			full_length += v.second.length;

		if (full_length > k)
			break;
	}

	print_sub("[", n - 1, k - n + 1);
}

Problem Statements