Submission #1309556

#	Time	Username	Problem	Language	Result	Execution time	Memory
1309556		nanaseyuzuki	Construction of Highway (JOI18_construction)	C++20	100 / 100	392 ms	21072 KiB

construction

#include <bits/stdc++.h>
// Kazusa_Megumi
#define ll long long
#define fi first
#define se second
#define pii pair<int, int>
#define all(a) a.begin(), a.end()
using namespace std;

#ifdef LOCAL
#include "C:\Users\Dell\Downloads\template\template\icpc-notebook\Utilities\debug.h"
#else
#define debug(...) 42
#endif

const int mn = 1e5 + 5, mod = 1e9 + 7, inf = 2e9;

int n, c[mn], d[mn], sz[mn], heavy[mn], par[mn], head[mn], st[mn], ft[mn], euler[mn], timer_dfs = 0;
vector <int> a[mn];
array <int, 2> e[mn];

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

void decompose(int u, int h){
	st[u] = ++ timer_dfs; euler[timer_dfs] = u;
	head[u] = h;
	if(heavy[u]) decompose(heavy[u], h);
	for(auto v : a[u]){
		if(v != heavy[u]) decompose(v, v);
	}
	ft[u] = timer_dfs;
}

int bit[mn];

void add(int u, int val){
	for(; u <= n; u += (u & -u)) bit[u] += val;
}

int get(int u){
	int res = 0;
	for(; u; u -= (u & -u)) res += bit[u];
	return res;
}

int st1[mn << 2], lz[mn << 2];

int merge(int a, int b){ return (a == b ? a : 0); };

void down(int id, int l, int r){
	if(l != r && lz[id]){
		st1[id << 1] = st1[id << 1 | 1] = lz[id << 1] = lz[id << 1 | 1] = lz[id];
		lz[id] = 0;
	}
}

void update(int id, int l, int r, int u, int v, int val){
	if(l > v || r < u) return;
	if(l >= u && r <= v){
		st1[id] = val;
		lz[id] = val;
		return;
	}
	down(id, l, r);
	int mid = (l + r) >> 1;
	update(id << 1, l, mid, u, v, val); update(id << 1 | 1, mid + 1, r, u, v, val);
	st1[id] = merge(st1[id << 1], st1[id << 1 | 1]);
}

stack <pii> roll_back;

int dfs(int id, int l, int r, int u, int v){
	if(l > v || r < u) return 0;
	if(l >= u && r <= v && st1[id]){
		roll_back.push({r - l + 1, st1[id]});
		add(st1[id], r - l + 1);
		return (r - l + 1) * get(st1[id] - 1);
	}
	int mid = (l + r) >> 1;
	down(id, l, r);
	return dfs(id << 1 | 1, mid + 1, r, u, v) + dfs(id << 1, l, mid, u, v);
}

void solve() {
    cin >> n;
    vector <int> comp;
    for(int i = 1; i <= n; i++){
    	cin >> c[i]; comp.push_back(c[i]);
    }
    sort(all(comp)); comp.erase(unique(all(comp)), comp.end());
    for(int i = 1; i <= n; i++) c[i] = lower_bound(all(comp), c[i]) - comp.begin() + 1;
    for(int i = 1; i < n; i++){
    	cin >> e[i][0] >> e[i][1];
    	auto [u, v] = e[i];
    	a[u].push_back(v);
    }
    dfs(1); decompose(1, 1);
    update(1, 1, n, st[1], st[1], c[1]);
    for(int i = 1; i < n; i++){
    	auto [u, v] = e[i];
    	update(1, 1, n, st[v], st[v], c[v]);
    	int res = 0;
    	while(u){
    		res += dfs(1, 1, n, st[head[u]], st[u]);
    		update(1, 1, n, st[head[u]], st[u], c[v]);
    		u = par[head[u]];
    	}
    	while(roll_back.size()){
    		auto [len, val] = roll_back.top();
    		add(val, - len); roll_back.pop();
    	}
    	cout << res << '\n';
    }
}

main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    if (fopen("Kazuki.INP", "r")) {
        freopen("Kazuki.INP", "r", stdin);
        freopen("Kazuki.OUT", "w", stdout);
    }
    int t = 1;
    // cin >> t;
    while (t--) solve();
    return 0;
}

Compilation message (stderr)

construction.cpp:125:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
  125 | main() {
      | ^~~~
construction.cpp: In function 'int main()':
construction.cpp:129:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  129 |         freopen("Kazuki.INP", "r", stdin);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~
construction.cpp:130:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  130 |         freopen("Kazuki.OUT", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~

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