Submission #1261897

#	Time	Username	Problem	Language	Result	Execution time	Memory
1261897		mingga	Construction of Highway (JOI18_construction)	C++20	100 / 100	164 ms	19388 KiB

construction

// Author: caption_mingle
#include "bits/stdc++.h"

using namespace std;

#define ln "\n"
#define pb push_back
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define sz(x) ((int)(x).size())
#define ll long long
const int mod = 1e9 + 7;
const int inf = 2e9;
const int N = 1e5 + 7;
int n, par[N], h[N], heavy[N], sz[N], nxt[N], a[N];
vector<int> g[N];
vector<pair<int, int>> que;
vector<pair<int, int>> chain[N];

struct BIT {
	vector<ll> bit;
	int n;
	BIT() {}
	BIT(int n) : n(n) {
		bit.resize(n + 1, 0);
	}
	void update(int u, int x) {
		for(; u <= n; u += (u & -u)) bit[u] += x;
	}
	ll get(int u) {
		ll ans = 0;
		for(; u > 0; u -= (u & -u)) ans += bit[u];
		return ans;
	}
};



void dfs(int u) {
	sz[u] = 1;
	for(int v : g[u]) {
		h[v] = h[u] + 1;
		par[v] = u;
		dfs(v);
		sz[u] += sz[v];
		if(sz[v] > sz[heavy[u]]) heavy[u] = v;
	}
}

void hld(int u) {
	for(int v : g[u]) {
		nxt[v] = (v == heavy[u] ? nxt[u] : v);
		hld(v);
	}
}

ll query(int u) {
	vector<pair<int, int>> dak;
	vector<int> val;
	while(u) {
		int pu = nxt[u];
		int mx = h[u] - h[pu] + 1;
		vector<pair<int, int>> tmp;
		for(int i = sz(chain[pu]) - 1; i >= 0; i--) {
			val.pb(chain[pu][i].fi);
			// cerr << chain[pu][i].fi << ' ' << chain[pu][i].se << ln;
			if(chain[pu][i].se <= mx) {
				mx -= chain[pu][i].se;
				tmp.pb(chain[pu][i]);
			} else {
				tmp.pb({chain[pu][i].fi, mx});
				break;
			}
		}
		// cerr << u << ' ' << pu << ' ' << par[pu] << ln;
		u = par[pu];
		for(int i = sz(tmp) - 1; i >= 0; i--) dak.pb(tmp[i]);
	}
	sort(all(val));
	val.erase(unique(all(val)), val.end());
	BIT bit(sz(val));
	ll ans = 0;
	for(auto [num, cnt] : dak) {
		// cerr << num << ' ' << cnt << ln;
		num = upper_bound(all(val), num) - val.begin();
		ans += bit.get(num - 1) * cnt;
		bit.update(num, cnt);
	}
	return ans;
}

void update(int u, int col) {
	while(u) {
		int pu = nxt[u];
		int mx = h[u] - h[pu] + 1;
		while(sz(chain[pu])) {
			if(chain[pu].back().se <= mx) {
				mx -= chain[pu].back().se;
				chain[pu].pop_back();
			} else {
				chain[pu].back().se -= mx;
				break;
			}
		}
		chain[pu].pb({col, h[u] - h[pu] + 1});
		// cerr << "CHAIN " << pu << ln;
		// for(auto [v, cnt] : chain[pu]) cerr << v << ' ' << cnt << ln;
		u = par[pu];
	}
}

signed main() {
	cin.tie(0) -> sync_with_stdio(0);
	#define task ""
	if(fopen(task ".INP", "r")) {
		freopen(task ".INP", "r", stdin);
		freopen(task ".OUT", "w", stdout);
	}
	cin >> n;
	for(int i = 1; i <= n; i++) cin >> a[i];
	for(int i = 1; i < n; i++) {
		int u, v; cin >> u >> v;
		que.pb({u, v});
		g[u].pb(v);
	}
	nxt[1] = 1;
	dfs(1);
	hld(1);
	chain[1].pb({a[1], 1});
	for(auto [u, v] : que) {
		cout << query(u) << ln;
		update(v, a[v]);
	}
    cerr << "\nTime: " << clock() * 1000 / CLOCKS_PER_SEC;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:117:24: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  117 |                 freopen(task ".INP", "r", stdin);
      |                 ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
construction.cpp:118:24: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  118 |                 freopen(task ".OUT", "w", stdout);
      |                 ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84