Submission #337779

#	Time	Username	Problem	Language	Result	Execution time	Memory
337779		er888kh	Construction of Highway (JOI18_construction)	C++17	100 / 100	1049 ms	25564 KiB

construction

#include <bits/stdc++.h>

using namespace std;

#define ALL(iter) (iter).begin(),(iter).end()

typedef int64_t s64;
typedef pair<int, int> pii;
typedef pair<s64, s64> pll;

#define F first
#define S second

#define MAXN 100005
#define SST 524288
#define HST 262144
#define LG 18

//dfs stuff
int bg[MAXN];
int ed[MAXN];
int hei[MAXN];
int spt[MAXN][LG];
int ct = 1;
//dfs stuff

//input stuff
int n;
int cs[MAXN];
int cnum[MAXN];
vector<pii> edgs;
vector<int> adj[MAXN];
//input stuff

//seg tree
int st[SST];
//seg tree

//fenwick tree
int bit[MAXN];
//fenwick tree

void update(int pos, int num){
	pos += HST;
	st[pos] = num;
	do{
		pos >>= 1;
		st[pos] = num;
	}while(pos > 1);
}

int query(int l, int r){
	l += HST;
	r += HST;
	int ans = 0;
	while(l <= r){
		if(l & 1) ans = max(ans, st[l++]);
		if(r+1&1) ans = max(ans, st[r--]);
		l >>= 1; r >>= 1;
	}
	return ans;
}

void dfs(int st, int pr, int h){
	bg[st] = ct++;
	hei[st] = h;
	spt[st][0] = pr;
	for(int i = 1; i < LG; i++){
		spt[st][i] = spt[spt[st][i-1]][i-1];
	}
	for(auto x:adj[st]){
		if(x == pr){
			continue;
		}
		dfs(x, st, h+1);
	}
	ed[st] = ct++;
}

int get_bit(int pos){
	int ans = 0;
	for(; pos; pos -= pos & -pos){
		ans += bit[pos];
	}
	return ans;
}

void set_bit(int pos, int val){
	for(; pos < MAXN; pos += pos & -pos){
		bit[pos] += val;
	}
}

void clear_bit(int pos){
	for(; pos < MAXN; pos += pos & -pos){
		bit[pos] = 0;
	}
}

int get_col_num(int v){
	return query(bg[v], ed[v]);
}

int get_col(int v){
	return cnum[get_col_num(v)];
}

int get_par(int v, int h){ //get v's h'th parent
	for(int i = 0; i < LG; i++){
		if(h & (1 << i)){
			v = spt[v][i];
		}
	}
	return v;
}

pii get_cols(int v){
	int l = 0;
	int r = hei[v]-1;
	int ccol = get_col_num(v);
	while(l < r){
		int mid = (l + r + 1) >> 1;
		int u = get_par(v, mid);
		int ncol = get_col_num(u);
		if(ncol == ccol){
			l = mid;
		} else {
			r = mid-1;
		}
	}
	//fprintf(stderr, "%d: {%d, %d}\n", v, ccol, l+1);
	return {ccol, l+1}; //color number of v, how many above v have the same color
	//real color of v is cnum[ccol]
}

int main(){
	scanf("%d", &n);
	{
		vector<int> ccs;
		for(int i = 1; i <= n; i++){
			scanf("%d", &cs[i]);
			ccs.push_back(cs[i]);
		}
		sort(ALL(ccs));
		ccs.erase(unique(ALL(ccs)), ccs.end());
		for(int i = 1; i <= n; i++){
			cs[i] = lower_bound(ALL(ccs), cs[i]) - ccs.begin() + 1;
		}
	}
	cnum[0] = cs[1];
	for(int i = 1; i < n; i++){
		int a, b;
		scanf("%d%d", &a, &b);
		cnum[i] = cs[b];
		adj[a].push_back(b);
		adj[b].push_back(a);
		edgs.push_back({a, b});
	}
	dfs(1, 0, 1);
	//update(1, 0);
	{
		int ce = 1;
		for(auto x:edgs){
			int a = x.F;
			int b = x.S;
			int cv = a;
			s64 tans = 0;
			vector<int> toclear;
			while(cv){
				//fprintf(stderr, "%d\n", cv);
				auto p = get_cols(cv);
				int rc = cnum[p.F];
				toclear.push_back(rc);
				cv = get_par(cv, p.S);
				s64 ct = p.S;
				tans += ct * 1LL * get_bit(rc-1);
				set_bit(rc, ct);
			}
			printf("%lld\n", tans);
			for(auto x:toclear){
				clear_bit(x);
			}
			update(bg[b], ce++); 
		}
	}
	return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int query(int, int)':
construction.cpp:58:7: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   58 |   if(r+1&1) ans = max(ans, st[r--]);
      |      ~^~
construction.cpp: In function 'int main()':
construction.cpp:179:15: warning: format '%lld' expects argument of type 'long long int', but argument 2 has type 's64' {aka 'long int'} [-Wformat=]
  179 |    printf("%lld\n", tans);
      |            ~~~^     ~~~~
      |               |     |
      |               |     s64 {aka long int}
      |               long long int
      |            %ld
construction.cpp:137:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  137 |  scanf("%d", &n);
      |  ~~~~~^~~~~~~~~~
construction.cpp:141:9: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  141 |    scanf("%d", &cs[i]);
      |    ~~~~~^~~~~~~~~~~~~~
construction.cpp:153:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  153 |   scanf("%d%d", &a, &b);
      |   ~~~~~^~~~~~~~~~~~~~~~

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