Submission #220216

#	Time	Username	Problem	Language	Result	Execution time	Memory
220216		atoiz	Harvest (JOI20_harvest)	C++14	100 / 100	776 ms	117560 KiB

harvest

/*input
4 4 12 1
3 2 7 5 
0 1 9 8 
1
2 6

*/

#include <iostream>
#include <vector>
#include <algorithm> 
#include <cassert>
#include <ctime>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>

using namespace std;
using namespace __gnu_pbds;

const int INF = 2 * 1000 * 1000 * 1000 + 9;

using pii = pair<int64_t, int>;	

using ordered_set = tree<pii, null_type, less<pii>, rb_tree_tag, tree_order_statistics_node_update>;
int ordered_counter = 0;
void set_insert(ordered_set &os, int64_t x)
{ os.insert(make_pair(x, ++ordered_counter)); }
int get_order(ordered_set &os, int64_t x)
{ return (int) os.order_of_key(make_pair(x, 0)); }

#define for_range(var, from, to) for (auto var = from; var < to; ++var)
#define for_each(var, list) for (auto &var : list)
#define all(list) begin(list), end(list)

int num_apples, num_employees, num_queries;
vector<int> pos_apples, pos_employees;
vector<vector<pii>> queries;
vector<int64_t> answers;
int perimeter, revive;

vector<pii> next_employee;
vector<vector<pii>> graph;
vector<vector<pii>> harvest_apples;
vector<bool> in_cycle;

void init_relations()
{
	vector<pii> timestamps_employee;
	timestamps_employee.reserve(2 * num_employees);
	for_range(employee, 0, num_employees) {
		timestamps_employee.emplace_back(pos_employees[employee], employee);
		timestamps_employee.emplace_back(pos_employees[employee] + perimeter, employee);
	}
	sort(all(timestamps_employee));

	// init next_employee
	next_employee.resize(num_employees);
	graph.assign(num_employees, vector<pii>(0));

	int revive_dist = revive % perimeter;
	int64_t revive_add = revive - revive_dist;
	for_range(time_id, num_employees, 2 * num_employees) {
		auto cur = timestamps_employee[time_id];
		int pos = cur.first, employee = cur.second;
		auto it = prev(upper_bound(all(timestamps_employee), pii(pos - revive_dist, INF)));

		int64_t wait_time = pos - it->first + revive_add;
		assert(wait_time >= revive);
		int next_id = it->second;
		next_employee[employee] = make_pair(wait_time, next_id);
		graph[next_id].emplace_back(wait_time, employee);

		// cout << "next " << employee << ": " << wait_time << ' ' << next_id << endl;
	}

	// init harvest_apples
	harvest_apples.assign(num_employees, vector<pii>(0));
	for_range(apple, 0, num_apples) {
		int pos = pos_apples[apple] + perimeter;
		auto it = prev(upper_bound(all(timestamps_employee), pii(pos, INF)));

		int wait_time = pos - it->first;
		int employee = it->second;
		harvest_apples[employee].emplace_back(wait_time, apple);

		// cout << "harvest " << employee << ": " << wait_time << endl;
	}
}

void find_cycles()
{
	in_cycle.assign(num_employees, false);
	vector<int> color(num_employees, 0);

	for_range(start_employee, 0, num_employees) {
		for (int employee = start_employee; color[employee] != 2; employee = next_employee[employee].second) {
			++color[employee];
			if (color[employee] == 2) in_cycle[employee] = true;
		}

		for (int employee = start_employee; color[employee] != 2; employee = next_employee[employee].second) {
			++color[employee];
		}
	}

	for_range(employee, 0, num_employees) {
		if (in_cycle[employee]) {
			for (auto it = graph[employee].begin(); it != graph[employee].end(); ++it) {
				if (in_cycle[it->second]) {
					graph[employee].erase(it);
					break;
				}
			}
		}
	}
}

vector<ordered_set> all_apples;
vector<int64_t> dist_inc;
void solve_tree(int cur_employee)
{
	for_each(edge, graph[cur_employee]) {
		int64_t dist = edge.first;
		int child_employee = edge.second;

		solve_tree(child_employee);
		dist_inc[child_employee] += dist;

		if (all_apples[cur_employee].size() < all_apples[child_employee].size()) {
			all_apples[cur_employee].swap(all_apples[child_employee]);
			swap(dist_inc[cur_employee], dist_inc[child_employee]);
		}

		for_each(apple, all_apples[child_employee]) {
			set_insert(all_apples[cur_employee], apple.first + dist_inc[child_employee] - dist_inc[cur_employee]);
		}
	}

	for_each(harvest, harvest_apples[cur_employee]) {
		set_insert(all_apples[cur_employee], harvest.first - dist_inc[cur_employee]);
	}

	for_each(query, queries[cur_employee]) {
		int64_t time = query.first;
		int id = query.second;

		answers[id] = get_order(all_apples[cur_employee], time - dist_inc[cur_employee]);
	}
}
void solve_all_trees()
{
	all_apples.assign(num_employees, ordered_set());
	dist_inc.assign(num_employees, 0);
	for_range(employee, 0, num_employees) {
		if (in_cycle[employee]) {
			solve_tree(employee);
		}
	}
}

void solve_cycle(vector<int> cycle)
{
	reverse(all(cycle));
	// for_each(employee, cycle) cout << employee << ' '; cout << endl;

	int cycle_size = (int) cycle.size();

	int64_t cycle_perimeter = 0;
	for_each(employee, cycle) cycle_perimeter += next_employee[employee].first;
	assert(cycle_perimeter % perimeter == 0);
	// cout << "perimeter " << cycle_perimeter << endl;

	vector<int64_t> relative_position(cycle_size, 0);
	for_range(cycle_id, 1, cycle_size) {
		relative_position[cycle_id] = relative_position[cycle_id - 1] + next_employee[cycle[cycle_id]].first;
	}

	// reset answers from tree solver
	for_each(employee, cycle) {
		for_each(query, queries[employee]) {
			answers[query.second] = 0;
		}
	}

	// prepare
	vector<int64_t> apple_times;
	vector<pair<int64_t, int>> cur_queries;

	for_range(cycle_id, 0, cycle_size) {
		int employee = cycle[cycle_id];
		int64_t pos_inc = relative_position[cycle_id];

		for_each(apple, all_apples[employee]) {
			apple_times.push_back(apple.first + dist_inc[employee] + pos_inc);
			// cout << apple.first + dist_inc[employee] << ' ' << pos_inc << endl;
		}

		for_each(query, queries[employee]) {
			cur_queries.emplace_back(query.first + pos_inc, query.second);
			// cout << query.first << ' ' << pos_inc << endl;
		}
	}
	sort(all(apple_times));
	sort(all(cur_queries));

	// remove unwanted count from previous (position-wise) employee
	{
		ordered_set cur_state;
		for_range(cycle_id, 0, cycle_size) {
			int employee = cycle[cycle_id];
			int64_t pos_inc = relative_position[cycle_id];

			// update answers
			for_each(query, queries[employee]) {
				answers[query.second] -= get_order(cur_state, query.first + pos_inc);
			}

			// update cur_state
			for_each(apple, all_apples[employee]) {
				set_insert(cur_state, apple.first + dist_inc[employee] + pos_inc);
			}
		}
	}

	// add extra count from x % C < T % C
	{
		ordered_set cur_state;
		auto apple_it = apple_times.begin();
		for_each(query, cur_queries) {
			int64_t query_time = query.first;
			for (; apple_it != apple_times.end() && *apple_it < query_time; ++apple_it) {
				set_insert(cur_state, *apple_it % cycle_perimeter);
			}

			answers[query.second] += get_order(cur_state, query_time % cycle_perimeter);
		}
	}

	// add main count
	{
		auto apple_it = apple_times.begin();
		int64_t apple_count = 0, excludant = 0;

		for_each(query, cur_queries) {
			int64_t query_time = query.first;
			for (; apple_it != apple_times.end() && *apple_it < query_time; ++apple_it) {
				++apple_count, excludant += *apple_it / cycle_perimeter;
			}

			answers[query.second] += apple_count * (query_time / cycle_perimeter) - excludant;
		}
	}
}
void solve_all_cycles()
{
	vector<bool> visited(num_employees, false);
	for_range(start_employee, 0, num_employees) {
		if ((not in_cycle[start_employee]) or visited[start_employee]) continue;
		vector<int> cur_cycle;
		for (int employee = start_employee; not visited[employee]; employee = next_employee[employee].second) {
			cur_cycle.push_back(employee);
			visited[employee] = true;
		}
		solve_cycle(cur_cycle);
	}
}

int32_t main()
{
	ios_base::sync_with_stdio(0); cin.tie(0);
	cin >> num_employees >> num_apples >> perimeter >> revive;
	pos_apples.resize(num_apples), pos_employees.resize(num_employees);
	for_each(pos, pos_employees) cin >> pos;
	for_each(pos, pos_apples) cin >> pos;

	cin >> num_queries;
	queries.assign(num_employees, vector<pii>(0));
	answers.resize(num_queries);
	for_range(query_id, 0, num_queries) {
		int employee;
		int64_t time;
		cin >> employee >> time;
		--employee;
		queries[employee].emplace_back(time + 1, query_id);
	}

	init_relations();
	find_cycles();
	solve_all_trees();
	solve_all_cycles();

	for_range(query_id, 0, num_queries) {
		cout << answers[query_id] << '\n';
	}
	// cout << 1.0 * clock() / CLOCKS_PER_SEC << endl;
	return 0;
}

• Subtask 1: 5 / 5
• Subtask 2: 20 / 20
• Subtask 3: 75 / 75