Submission #119920

#	Time	Username	Problem	Language	Result	Execution time	Memory
119920		arman_ferdous	Construction of Highway (JOI18_construction)	C++14	100 / 100	589 ms	95340 KiB

construction

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
const int N = 1e5+100;
int n, c[N];
vector<int> g[N];
struct edge{
	int u, v;
}e[N];

int par[N], sub[N], lev[N];
int subdfs(int u, int f, int l) {
	par[u] = f;
	sub[u] = 1;
	lev[u] = l;
	for(int v : g[u]) if(v != f) 
		sub[u] += subdfs(v, u, l + 1);
	return sub[u];
}

int head[N], cid[N], zap;
deque< pair<int,ll> > val[N];
void hld(int u, int f, int dad, int id) {
	head[u] = dad;
	cid[u] = id;

	if(!val[id].empty() && val[id].back().first == c[u])
		val[id].back().second++;
	else val[id].push_back({c[u], 1});

	int mx = -1, sc;
	for(int v : g[u]) if(v != f && sub[v] > mx)
		mx = sub[v], sc = v;
	if(mx == -1) return;
	hld(sc, u, dad, id);

	for(int v : g[u]) if(v != f && v != sc)
		hld(v, u, v, ++zap);
}

int bit[N];
deque< pair<int,ll> > tmp, full;

void upd(int pos, int x) {
	while(pos < N) bit[pos] += x, pos += pos&-pos;
}
ll get(int pos, int sum = 0) {
	while(pos > 0) sum += bit[pos], pos -= pos&-pos;
	return sum;
}

ll calc(int u, int newc) {
	full.clear();
	while(u != -1) {
		tmp.clear();
		int take = lev[u] - lev[head[u]] + 1, id = cid[u];
		while(take) {
			if(take - val[id][0].second <= 0) {
				tmp.push_front({val[id][0].first, take});
				val[id][0].second -= take;
				if(val[id][0].second == 0) val[id].pop_front();
				break;
			}
			else {
				tmp.push_front(val[id][0]);
				take -= val[id][0].second;
				val[id].pop_front();
			}
		}
		if(!val[id].empty() && val[id][0].first == newc) 
			val[id][0].second += lev[u] - lev[head[u]] + 1;
		else val[id].push_front({newc, lev[u] - lev[head[u]] + 1});
		u = par[head[u]];
		while(!tmp.empty()) {
			full.push_front(tmp.front());
			tmp.pop_front();
		}
	}
	int sz = full.size(); ll ret = 0;
	for(int i = sz - 1; i >= 0; i--) {
		ret += get(full[i].first - 1) * full[i].second;
		upd(full[i].first, full[i].second);
	}
	for(int i = 0; i < sz; i++) 
		upd(full[i].first, -full[i].second);
	return ret;
}

vector<int> v;
map<int,int> mp;
int main() {
	scanf("%d", &n);
	for(int i = 1; i <= n; i++) {
		scanf("%d", &c[i]);
		v.push_back(c[i]);
	} sort(v.begin(), v.end());
	v.erase(unique(v.begin(),v.end()),v.end());
	int cur = 1;
	for(auto it : v) mp[it] = cur++;
	for(int i = 1; i <= n; i++) c[i] = mp[c[i]];

	for(int i = 1; i < n; i++) {
		scanf("%d %d", &e[i].u, &e[i].v);
		g[e[i].u].push_back(e[i].v);
		g[e[i].v].push_back(e[i].u);
	}
	subdfs(1, -1, 1);
	hld(1, -1, 1, ++zap);

	for(int i = 1; i < n; i++)
		printf("%lld\n", calc(e[i].u, c[e[i].v]));
	return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:93:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d", &n);
  ~~~~~^~~~~~~~~~
construction.cpp:95:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d", &c[i]);
   ~~~~~^~~~~~~~~~~~~
construction.cpp:104:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d", &e[i].u, &e[i].v);
   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~

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