Submission #1180802

#	Time	Username	Problem	Language	Result	Execution time	Memory
1180802		jbareja	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	353 ms	53884 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;

struct SegTree {
	vector<int> values;
	int base;

	void Init(int n) {
		base = 1;
		while (base <= n) base *= 2;
		values.assign(base * 2, 0);
	}

	void Add(int l, int r, int val) {
		l += base - 1;
		r += base + 1;
		while (l/2 != r/2) {
			if (l % 2 == 0) values[l+1] += val;
			if (r % 2 != 0) values[r-1] += val;
			l /= 2, r /= 2;
		}
	}

	int Sum(int x) {
		x += base;
		int ans = 0;
		while (x) {
			ans += values[x];
			x /= 2;
		}
		return ans;
	}
};

struct CntUniques {
	int cnt_uniques = 0;
	vector<int> cnt_values;

	void Init(int n) {
		cnt_values.assign(n+1, 0);
	}

	void Insert(int x) {
		if (!cnt_values[x]) cnt_uniques++;
		cnt_values[x]++;
	}

	void Erase(int x) {
		cnt_values[x]--;
		if (!cnt_values[x]) cnt_uniques--;
	}

	int Uniques() {
		return cnt_uniques;
	}
};

int N; // liczba wierzchołków
int M; // maksymalny numer specjalności

vector<vector<int>> graph;
vector<int> speciality;

vector<int> depth;
vector<int> answer;
int max_depth = 0;
SegTree tree;
vector<int> height;
CntUniques U;

int dfs_farthest_from(int v, int p=0) {
	if (p == 0) depth[v] = 1;
	int ans = v;
	for (int u: graph[v]) if (u != p) {
		depth[u] = depth[v] + 1;
		max_depth = max(depth[u], max_depth);
		int temp = dfs_farthest_from(u, v);
		if (depth[temp] > depth[ans]) ans = temp;
	}
	return ans;
}

void dfs_subtree_height(int v, int p=0) {
	if (p == 0) depth[v] = 1;
	height[v] = 1;
	for (int u: graph[v]) if (u != p) {
		depth[u] = depth[v] + 1;
		dfs_subtree_height(u, v);
		height[v] = max(height[v], height[u] + 1);
	}
}

vector<int> path = { 0 };

void dfs_count_answer(int v, int p=0, int upper=0, int bottom=0) {
	path.push_back(v);
	vector<int> unlocked;
	for (int u: graph[v]) if (u != p) {
		int upper_border = max(0, depth[v]-height[u]);
		int bottom_border = depth[p];
		tree.Add(upper_border, bottom_border, 1);
	}
	if (depth[v] > 1) {
		tree.Add(upper, bottom, -1);
		for (int i=max(1,upper); i<upper+2 && i<depth[v]; i++) {
			if (tree.Sum(i) == 0) {
				unlocked.push_back(speciality[path[i]]);
				U.Insert(speciality[path[i]]);
			}
		}
	}
	answer[v] = max(answer[v], U.Uniques());
	for (int u: graph[v]) if (u != p) {
		int upper_border = max(0, depth[v]-height[u]);
		int bottom_border = depth[p];
		dfs_count_answer(u, v, upper_border, bottom_border);
	}

	path.pop_back();
	for (int u: graph[v]) if (u != p) {
		int upper_border = max(0, depth[v]-height[u]);
		int bottom_border = depth[p];
		tree.Add(upper_border, bottom_border, -1);
	}
	if (depth[v] > 1) {
		tree.Add(upper, bottom, 1);
		for (int i: unlocked) U.Erase(i);
	}
}

void count_from(int v_start) {
	dfs_subtree_height(v_start);
	depth[0] = 0;
	dfs_count_answer(v_start);
}

int main() {
	ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);

	cin >> N >> M;

	graph.assign(N+1, vector<int>());
	speciality.assign(N+1, 0);
	depth.assign(N+1, 0);
	answer.assign(N+1, 0);
	height.assign(N+1, 0);
	U.Init(M);

	for (int i=0; i<N-1; i++) {
		int A, B;
		cin >> A >> B;
		graph[A].push_back(B);
		graph[B].push_back(A);
	}

	for (int i=1; i<=N; i++) cin >> speciality[i];

	int vA = dfs_farthest_from(1);
	int vB = dfs_farthest_from(vA);

	tree.Init(max_depth + 1);

	count_from(vA);
	count_from(vB);

	for (int i=1; i<=N; i++) cout << answer[i] << '\n';
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32