Submission #490912

#	Time	Username	Problem	Language	Result	Execution time	Memory
490912		nafis_shifat	Construction of Highway (JOI18_construction)	C++17	100 / 100	288 ms	22168 KiB

construction

#include<bits/stdc++.h>
#define ll long long
#define pii pair<int,int>
using namespace std;
const int mxn=1e5+5;
const int inf=1e9;
int n, c[mxn];
vector<int> adj[mxn];
int U[mxn], V[mxn];

int chainNo = 1;
int ptr = 1;

int parent[mxn];


int chainHead[mxn];
int chainInd[mxn];
int posInBase[mxn];
int subsize[mxn];
int chainLen[mxn] = {};

vector<pii> con[mxn];



struct BIT {
	int bit[mxn] = {};
	vector<pii> ups;


	void update(int p,int v, bool rem = true) {
		if(p <= 0) return;
		if(rem) ups.emplace_back(p, v);
		for(;p<mxn;p+=p&-p) bit[p] += v;
	}
    int query(int p) {
    	int r=0;
    	for(;p>0;p-=p&-p) r += bit[p];
        return r;
    }

    void clear() {
    	while(!ups.empty()) {
    		update(ups.back().first, -ups.back().second, false);
    		ups.pop_back();
    	}
    }
} bt;

void dfs(int u, int prev) {
	subsize[u] = 1;
	for(int v : adj[u]) {
		if(v == prev) continue;

		parent[v] = u;
		dfs(v, u);
		subsize[u] += subsize[v];
	}
}



void HLD(int curNode, int prev) {

	if(chainHead[chainNo] == -1) {
		chainHead[chainNo] = curNode; 

	}

	chainLen[chainNo]++;
	chainInd[curNode] = chainNo;
	posInBase[curNode] = ptr++; 
	int sc = -1;
	for(int i = 0; i < adj[curNode].size(); i++) if(adj[curNode][i] != prev){
		if(sc == -1 || subsize[sc] < subsize[adj[curNode][i]]) {
			sc = adj[curNode][i];
		}
	}

	if(sc != -1) {
		HLD(sc, curNode);
	}

	for(int i=0; i < adj[curNode].size(); i++) if(adj[curNode][i] != prev) {
		if(sc != adj[curNode][i]) {
			chainNo++;
			HLD(adj[curNode][i], curNode);
		}
	}
}



ll get(int u, int v, int x) {
	int dn = 0;

	int len = posInBase[v] - posInBase[u] + 1;

	int id = chainInd[u];
	ll ans = 0;
	vector<pii> prc;
	while(dn < len) {
		if(dn + con[id].back().first > len) {
			pii lst = con[id].back();
			con[id].pop_back();
			int dif = len - dn;

			prc.emplace_back(dif, lst.second);


			dn = len;

			con[id].push_back({lst.first - dif, lst.second});
			break;
		}

		assert(!con[id].empty());

		pii lst = con[id].back();
		con[id].pop_back();
		


		prc.emplace_back(lst.first, lst.second);

		dn += lst.first;
	}

	con[id].push_back({len, x});


	while(!prc.empty()) {
		ans += 1ll * prc.back().first * bt.query(prc.back().second - 1);
		bt.update(prc.back().second, prc.back().first);
		prc.pop_back();
	}

	return ans;
}


ll query_up(int u, int x) {
	if(u == 1) return 0; 
	int uchain, ans = 0;
	while(u > 0) {
		uchain = chainInd[u];
		if(uchain == 1) {
			ans += get(1, u, x);
			break;
		}

		ans += get(chainHead[uchain], u, x);
		
		u = chainHead[uchain]; 
		u = parent[u];
	}
	return ans;
}


int main() {

	memset(chainHead, -1, sizeof chainHead);
	cin >> n;

	vector<int> vv;

	for(int i = 1; i <= n; i++) {
		scanf("%d", &c[i]);
		vv.push_back(c[i]);
	}

	sort(vv.begin(), vv.end());
	vv.erase(unique(vv.begin(), vv.end()), vv.end());

	for(int i = 1; i <= n; i++) {
		c[i] = lower_bound(vv.begin(), vv.end(), c[i]) - vv.begin() + 1;
	}

	for(int i = 1; i < n; i++) {
		scanf("%d%d", &U[i], &V[i]);
		adj[U[i]].push_back(V[i]);
		adj[V[i]].push_back(U[i]);
	}

	dfs(1, -1);


	HLD(1, -1);

	for(int i = 1; i <= n && chainHead[i] != -1; i++) {
		con[i].emplace_back(chainLen[i], -1);
	}


	for(int i = 1; i < n; i++) {
		cout<<query_up(V[i], c[V[i]])<<"\n";
		bt.clear();
	}
	
}

Compilation message (stderr)

construction.cpp: In function 'void HLD(int, int)':
construction.cpp:75:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   75 |  for(int i = 0; i < adj[curNode].size(); i++) if(adj[curNode][i] != prev){
      |                 ~~^~~~~~~~~~~~~~~~~~~~~
construction.cpp:85:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   85 |  for(int i=0; i < adj[curNode].size(); i++) if(adj[curNode][i] != prev) {
      |               ~~^~~~~~~~~~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:170:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  170 |   scanf("%d", &c[i]);
      |   ~~~~~^~~~~~~~~~~~~
construction.cpp:182:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  182 |   scanf("%d%d", &U[i], &V[i]);
      |   ~~~~~^~~~~~~~~~~~~~~~~~~~~~

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