Submission #1310243

#TimeUsernameProblemLanguageResultExecution timeMemory
1310243mat_jurConstruction of Highway (JOI18_construction)C++20
0 / 100
1 ms568 KiB
#include "bits/stdc++.h"
using namespace std;
#ifdef DEBUG
auto&operator<<(auto &o, pair<auto, auto> p) {o << "(" << p.first << ", " << p.second << ")"; return o;}
auto operator<<(auto&o,auto x)->decltype(x.end(),o){o<<"{"; for(auto e : x) o<<e<<", "; return o<<"}";}
#define debug(X) cerr << "["#X"]: " << X << '\n';
#else 
#define cerr if(0)cout
#define debug(X) ;
#endif
using ll = long long;
#define all(v) (v).begin(), (v).end()
#define ssize(x) int(x.size())
#define fi first
#define se second
#define mp make_pair
#define eb emplace_back

struct segtree {
	int base;
	vector<int> t;

	segtree() {}
	void resize(int n) {
		base = 1;
		while (base < n) base *= 2;
		t.resize(2 * base);
	}

	void update(int v, int x) {
		v += base;
		while (v)
			t[v] += x, v /= 2;
	}

	int query(int l, int r) {
		int res = 0;
		l += base - 1; r += base + 1;
		while (l/2 != r/2) {
			if (l%2 == 0) res += t[l+1];
			if (r%2 == 1) res += t[r-1];
			l /= 2; r /= 2;
		}

		return res;
	}
};

struct HLD {
	int n;
	vector<vector<int>> g;
	vector<int> d, path, parent, sz, t;
	vector<set<pair<pair<int, int>, int>>> s;
	segtree T;

	HLD(const vector<vector<int>> &graph): n(ssize(graph)), g(graph) {
		d = path = parent = sz = t = vector<int>(n, -1);
		T.resize(n);
		s.resize(n);

		d[0] = 0;
		dfs(0, -1);
		path[0] = 0;
		make_hld(0, -1);

		debug(parent);
		debug(path);
		debug(d);
	}

	void dfs(int v, int p) {
		parent[v] = p;
		sz[v] = 1;
		t[v] = -1;
		for (int u : g[v]) {
			if (u == p) continue;
			d[u] = d[v] + 1;
			dfs(u, v);
			sz[v] += sz[u];
			t[v] = t[v] == -1 ? u : sz[t[v]] > sz[u] ? t[v] : u;
		}
	}

	void make_hld(int v, int p) {
		if (t[v] != -1) {
			path[t[v]] = path[v];
			make_hld(t[v], v);
		}

		for (int u : g[v]) {
			if (u == p || u == t[v]) continue;
			path[u] = u;
			make_hld(u, v);
		}
	}

	void update(int v, int c) {
		while (v != -1) {
			int p_start = path[v];
			pair<pair<int, int>, int> last;
			while (ssize(s[p_start]) && d[s[p_start].begin()->fi.fi] <= d[v]) {
				last = *s[p_start].begin();
				s[p_start].erase(s[p_start].begin());
			}

			if (d[last.fi.se] > d[v]) {
				last.fi.fi = d[v] + 1;
				s[p_start].insert(last);
			}
			s[p_start].insert({{d[p_start], d[v]}, c});

			v = parent[p_start];
		}
	}

	ll query(int v) {
		debug(s);
		vector<pair<int, int>> p;
		while (v != -1) {
			int p_start = path[v];
			auto it = s[p_start].begin();
			vector<pair<int, int>> p1;
			while (it != s[p_start].end() && d[it->fi.fi] <= d[v]) {
				p1.eb(min(d[v], it->fi.se) - it->fi.fi + 1, it->se);
				++it;
			}

			for (int i = ssize(p1) - 1; i >= 0; --i) p.eb(p1[i]);

			v = parent[p_start];
		}

		ll res = 0;
		
		debug(p)

		for (auto [cnt, x] : p) {
			res += ll(cnt) * T.query(0, x-1);
			T.update(x, cnt);
		}

		for (auto [cnt, x] : p)
			T.update(x, -cnt);

		return res;
	}
};

int main() {
	ios_base::sync_with_stdio(false); cin.tie(nullptr);

	int n;
	cin >> n;
	vector<int> c(n);
	map<int, int> remap;
	for (int &x : c) {
		cin >> x;
		remap[x];
	}

	int cnt = 0;
	for (auto &e : remap)
		e.se = cnt++;

	for (int &x : c) x = remap[x];

	vector<pair<int, int>> e(n-1);
	vector<vector<int>> g(n);
	for (auto &[a, b] : e) {
		cin >> a >> b;
		--a; --b;
		g[a].eb(b);
		g[b].eb(a);
	}

	HLD hld(g);
	hld.update(0, c[0]);

	for (auto [a, b] : e) {
		cout << hld.query(a) << '\n';
		hld.update(b, c[b]);
	}
	
	return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...