Submission #46774

#	Time	Username	Problem	Language	Result	Execution time	Memory
46774		kyleliu	Construction of Highway (JOI18_construction)	C++14	100 / 100	784 ms	178432 KiB

construction

#include <bits/stdc++.h> // PLEASE

using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair <int, int> pp;
#define MAXN 100005
#define MAXM 1005
#define MAXP 25
#define A first
#define B second
#define MP make_pair
#define PB push_back
#define PF push_front
const ll INF = 2e9+13;
const ll MOD = 1e9+7;
int N;
int C[MAXN];
vector <int> g[MAXN];
int in[MAXN], rin[MAXN], top[MAXN], out[MAXN], sz[MAXN], d[MAXN], par[MAXN];
int c = 1; 
void dfs1(int u)
{
	sz[u] = 1;
	for(auto& v : g[u]) {
		d[v] = d[u] + 1; dfs1(v); sz[u] += sz[v]; 
		if(sz[v] > sz[g[u][0]]) swap(g[u][0], v);
	}
}

void dfs2(int u)
{
	in[u] = c++;
	rin[in[u]] = u;
	for(auto v : g[u]) {
		if(v == g[u][0]) top[v] = top[u];
		else top[v] = v;
		dfs2(v);
	} out[u] = c;
}

struct D {
	int num, cnt, dep;
};

struct BIT { // 1-indexed
	ll T[MAXN];
	void init() { memset(T, 0, sizeof(T)); }
	void upd(int x, ll v) { while(x < MAXN) { T[x] += v; x += (x & -x); } }
	ll qry(int x) { ll ret = 0; while(x > 0) { ret += T[x]; x -= (x & -x); } return ret; }
	ll get(int a, int b) {
		if(a > b) return 0;
		return qry(b) - qry(a-1);
	}
} FEN;
deque <D> Q[MAXN];
vector <pp> ups;
vector <D> enc;
ll qry(int u)
{
	vector <int> todo;
	int tmp = u;
	while(top[tmp] != 1) {
		tmp = par[top[tmp]];
		todo.PB(top[tmp]);
	}
	reverse(todo.begin(), todo.end());
	ll ret = 0;
	ups.clear(); enc.clear();
	for(int i=0; i<todo.size(); i++) {
		int root = todo[i];
		int nxt = top[u]; if(i+1 < todo.size()) nxt = todo[i+1];
		int cn = 0;
		for(auto V : Q[root]) {
			cn += V.cnt; ll cv = V.cnt;
			if(cn >= d[nxt] - d[root]) cv -= (ll)(cn - (d[nxt] - d[root]));
			ret += FEN.get(V.num + 1, MAXN-1) * cv;
			FEN.upd(V.num, cv); ups.PB({V.num, cv});
			if(cn >= d[nxt] - d[root]) break;
		}
	}
	// now for the remaining
	int cn = 0;
	for(auto V : Q[top[u]]) {
		cn += V.cnt;
		ll cv = V.cnt;
		if(cn >= d[u] - d[top[u]] + 1) cv -= (ll)(cn - (d[u] - d[top[u]] + 1));
		ret += FEN.get(V.num + 1, MAXN-1) * cv;
		FEN.upd(V.num, cv); ups.PB({V.num, cv});
		if(cn >= d[u] - d[top[u]] + 1) break;
	}		
	for(auto up : ups) FEN.upd(up.A, -up.B);
	return ret;
}

void kill(int u, int col, int add)
{
	int t = top[u]; 
	int cn = 0; bool pushed = 0;
	while(Q[t].size() && (Q[t].front().dep <= d[u])) { 
		cn += Q[t].front().cnt;
		D tmp = Q[t].front();
		Q[t].pop_front(); 
		if(cn >= d[u] - d[t] + 1) { // split this one
			cn -= tmp.cnt;
			cn += d[u] - tmp.dep + 1;
			tmp.cnt -= d[u] - tmp.dep + 1;
			
			if(tmp.cnt) Q[t].PF({tmp.num, tmp.cnt, d[t] + cn});
			cn += add; pushed = 1;
			Q[t].PF({col, cn, d[t]}); cn = 0;
			break;
		}	
	}
	if(!pushed) Q[t].PF({col, cn+add, d[t]});
}
void upd(int u)
{
	int col = C[u];
	// update this chain
	kill(u, col, 1);
	// now kill all other chains
	while(top[u] != 1) {
		u = par[top[u]];
		kill(u, col, 0);
	}

	/*cout << endl;
	cout << "PRINTING" << endl;
	for(int i=1; i<=N; i++) {
		if(i == top[i]) {
			cout << "from node " << i << ":" << endl;
			for(auto V : Q[i]) cout << V.num << ' ' << V.cnt << ' ' << V.dep << endl;
		}
	}*/
}
int main(void)
{
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	cin >> N; vector <int> v;
	for(int i=1; i<=N; i++) { cin >> C[i]; v.PB(C[i]); }
	sort(v.begin(), v.end()); v.resize(unique(v.begin(), v.end()) - v.begin());
	for(int i=1; i<=N; i++) C[i] = lower_bound(v.begin(), v.end(), C[i]) - v.begin() + 1;
	vector <pp> ord;
	for(int i=0; i<N-1; i++) {
		int p, u; cin >> p >> u;
		par[u] = p;
		g[p].PB(u);
		ord.PB({p, u});
	} 
	par[1] = 1; top[1] = 1; dfs1(1); dfs2(1); Q[1].PF({C[1], 1, d[1]});
	FEN.init();
	
	for(auto q : ord) {
		int p = q.A; int u = q.B;
		cout << qry(p) << endl;
		upd(u);
	}
}

Compilation message (stderr)

construction.cpp: In function 'll qry(int)':
construction.cpp:70:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for(int i=0; i<todo.size(); i++) {
               ~^~~~~~~~~~~~
construction.cpp:72:28: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   int nxt = top[u]; if(i+1 < todo.size()) nxt = todo[i+1];
                        ~~~~^~~~~~~~~~~~~

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