제출 #239565

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
239565		jtnydv25	Construction of Highway (JOI18_construction)	C++17	100 / 100	1266 ms	66148 KiB

construction

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

#define ll long long
#define sd(x) scanf("%d", &(x))
#define pii pair<int, int>
#define F first
#define S second
#define all(c) ((c).begin()), ((c).end())
#define sz(x) ((int)(x).size())
#define ld long double

template<class T,class U>
ostream& operator<<(ostream& os,const pair<T,U>& p){
	os<<"("<<p.first<<", "<<p.second<<")";
	return os;
}

template<class T>
ostream& operator <<(ostream& os,const vector<T>& v){
	os<<"{";
	for(int i = 0;i < (int)v.size(); i++){
		if(i)os<<", ";
		os<<v[i];
	}
	os<<"}";
	return os;
}

#ifdef LOCAL
#define cerr cout
#else
#endif

#define TRACE

#ifdef TRACE
#define trace(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1>
void __f(const char* name, Arg1&& arg1){
	cerr << name << " : " << arg1 << std::endl;
}
template <typename Arg1, typename... Args>
void __f(const char* names, Arg1&& arg1, Args&&... args){
	const char* comma = strchr(names + 1, ',');cerr.write(names, comma - names) << " : " << arg1<<" | ";__f(comma+1, args...);
}
#else
#define trace(...)
#endif

// 0-indexed
// O(n) precomputation, O(1) query lca
// O(n logn) precomputation, O(1) query kth ancestor
struct tree {
	int n, block_size, num_blocks;
	vector<vector<int>> adj, val;
	vector<int> num, floorlog;
	tree(){}
	tree(int n) : n(n), adj(n){ }
	void add_edge(int s, int t) {
		adj[s].push_back(t);
		adj[t].push_back(s);
	}
	vector<int> pos, tour, depth, pos_end;
	vector<vector<int>> table;
	int argmin(int i, int j) { return depth[i] < depth[j] ? i : j; }
	// O(n) preprocessing, O(1) lca query
	void rootify(int r) {
		pos.resize(n);
		pos_end.resize(n);
		// euler tour
		function<void (int,int,int)> dfs = [&](int u, int p, int d) {
			pos[u] = pos_end[u] = depth.size();
			tour.push_back(u);
			depth.push_back(d);
			for (int v: adj[u]) {
				if (v != p) {
					dfs(v, u, d+1);
					pos_end[u] = (int)depth.size();
					tour.push_back(u);
					depth.push_back(d);
				}
			}
		}; 

		dfs(r, r, 0);

		floorlog.resize(tour.size() + 1);

		for(int i = 0; (1 << i) <= tour.size(); i++)
			for(int j = (1 << i); j < (1 << (i + 1)) && j <= tour.size(); j++)
				floorlog[j] = i;

		int logn = floorlog[tour.size()];
		block_size = logn / 2 + 1;
		num_blocks = ((int)tour.size() - 1) / block_size + 1;
		table.resize(logn+1, vector<int>(num_blocks));
		num.resize(num_blocks);
		val.resize(block_size + 1, vector<int>(1 << block_size));

		for(int i = 0; i < tour.size(); i += block_size){
			int mn = i;
			int v = 0;
			int block = i / block_size;
			int j = i;
			for(; j < tour.size() && j < i + block_size; j++){
				mn = argmin(mn, j);
				if(j != i){
					v = 2 * v + (depth[j] == depth[j - 1] + 1);
				}
			}
			table[0][block] = mn;
			num[block] = v << (block_size - j + i);
		}

		for (int h = 1; (1 << h) <= num_blocks; ++h) 
			for (int i = 0; i + (1<<h) <= num_blocks; ++i)
				table[h][i] = argmin(table[h - 1][i], table[h - 1][i+(1<<(h - 1))]);

		for(int i = 0; i < (1 << block_size); i++){
			int prefix = 0, mn = 0, curr = 0;
			for(int j = 0; j < block_size; j++){
				prefix += 2 * (i >> (block_size - 1 - j) & 1) - 1;
				if(mn > prefix){
					mn = prefix;
					curr = j + 1;
				}
				val[j + 1][i >> (block_size - 1 - j)] = curr;
			}
		}
	}

	inline int same_block_lca(int block, int i, int j){
		if(i == j) return i + block * block_size;
		int l = j - i;
		int mask = (num[block] >> (block_size - j - 1)) & ((1 << l) - 1);
		return block * block_size + i + val[l][mask];
	}

	int lca(int a, int b){
		a = pos[a]; b = pos[b];
		if(a > b) swap(a, b);
		int block_a = a / block_size;
		int block_b = b / block_size;
		int ind_a = a - block_a * block_size;
		int ind_b = b - block_b * block_size;
		if(block_a == block_b) return tour[same_block_lca(block_a, ind_a, ind_b)];
		int ans = argmin(same_block_lca(block_a, ind_a, block_size - 1), same_block_lca(block_b, 0, ind_b));
		if(block_b > block_a + 1){
			int h = floorlog[block_b - block_a - 1];
			int t = argmin(table[h][block_a + 1], table[h][block_b- (1<<h)]);
			ans = argmin(ans, t);
		}
		return tour[ans];
	}

	inline int dist(int i, int j){
		return depth[pos[i]] + depth[pos[j]] - 2 * depth[pos[lca(i, j)]];
	}

	inline int getDepth(int u){
		return depth[pos[u]];
	}

	inline bool isAncestor(int a, int b){
		return pos[b] >= pos[a] && pos[b] <= pos_end[a];
	}

	inline bool isOn(int c, int a, int b){
		return isAncestor(lca(a, b), c) && (isAncestor(c, a) || isAncestor(c, b));
	}
	inline bool pathsIntersect(int a, int b, int c, int d){
		int lab = lca(a, b), lcd = lca(c, d);
		return isOn(lab, c, d) || isOn(lcd, a, b);
	}
    vector<vector<int>> chains, par;
    vector<int> sizes, st, en, chain_index, maxDepth;

	void straighten(int r){
		const int logN = 20;
		st.resize(n);
		en.resize(n);
		sizes.resize(n);
		maxDepth.resize(n);
		par.assign(logN, vector<int>(n, 0));
		int timer = 0;
		function<void(int, int, int)> dfs = [&](int s, int p, int d){
			par[0][s] = p;
			st[s] = ++timer;
			maxDepth[s] = d;
			for(int v : adj[s]) if(v != p){
				dfs(v, s, d + 1);
				maxDepth[s] = max(maxDepth[s], maxDepth[v]);
			}
			en[s] = timer;
			sizes[s] = en[s] - st[s] + 1;
		};
		dfs(r, r, 0);
		for(int i = 1; i < logN; i++)
			for(int j = 0; j < n; j++) par[i][j] = par[i - 1][par[i - 1][j]];
	}
    // for kth ancestor in O(1)
    void hld(int r){
		straighten(r);
        function<void(int, int, vector<int> &)> dfs = [&](int s, int p, vector<int> & chain){
			int bigc = -1;
			for(int v : adj[s]) if(v != p){
				if(bigc == -1 || maxDepth[v] > maxDepth[bigc]) bigc = v;
			}
			for(int v : adj[s]) if(v != p){
				if(v == bigc){
					chain.push_back(v);
					dfs(v, s, chain);
				} else{
					vector<int> new_chain = {v};
					dfs(v, s, new_chain);
				}
			}
			if(bigc == -1) chains.push_back(chain);
        };
		vector<int> init_chain = {r};
		dfs(r, r, init_chain);
		chain_index.resize(n);
		int ind = 0;
		for(auto & it : chains){
			reverse(all(it));
			int l = sz(it);
			int v = it.back();
			it.resize(2 * l);
			for(int i = 0; i < l; i++){
				chain_index[it[i]] = ind;
				v = par[0][v];
				it[l + i] = v;
			}
			ind++;
		}
    }
	inline int msb(int x){
		return 31 - __builtin_clz(x);
	}
	int kthAncestor(int x, int k){
		if(k==0) return x;
		int w = msb(k);
		int v = par[w][x];
		k -= (1 << w);
		int ind = chain_index[v];
		int pos = getDepth(chains[ind][0]) - getDepth(v);
		return chains[ind][pos + k];
	}
	void input(){
        for(int i = 1; i < n - 1; i++){
            int x, y;
            scanf("%d %d", &x, &y);
            add_edge(x, y);
        }
    }
};

const int N = 100005;
const int logN = 20;

int arrival[N], par[logN][N], depth[N], A[N], B[N], C[N];
int st[N], en[N], timer;

vector<int> children[N];
void dfs(int s, int d){
	depth[s] = d;
	st[s] = timer++;
	for(int v : children[s]) dfs(v, d + 1);
	en[s] = timer - 1;
}

// 0-indexed
template<class T>
struct segtree{
	int n;
	vector<T> t;
	T def;
	inline T combine(T a, T b){
		return arrival[a] > arrival[b] ? a : b;
	}

	segtree(vector<T> & inp, T def = T()) : n(sz(inp)), def(def){
		t.resize(2 * n, def);
		for(int i = 0; i < n; i++) t[n + i] = inp[i];
		for(int i = n - 1; i > 0; --i) t[i] = combine(t[i<<1], t[i<<1|1]);
	}

	void modify(int p, T value) { // modify a[p] = value
		// value = combine(value, t[p + n]); // if a[p] = combine(a[p], value)
		for (t[p += n] = value; p >>= 1; ) t[p] = combine(t[p<<1], t[p<<1|1]);
	}

	T query(int l, int r) {  // compute on interval [l, r]
    r++;
		T resl = def, resr = def;
		for (l += n, r += n; l < r; l >>= 1, r >>= 1) {
			if (l&1) resl = combine(resl, t[l++]);
			if (r&1) resr = combine(t[--r], resr);
		}
		return combine(resl, resr);
	}
};

ll bit[N];
void add(int x, int y){
	x++;
	for(; x < N; x += x & -x) bit[x] += y;
}

ll get(int x){
	x++;
	ll ret = 0;
	for(; x; x -= x & -x) ret += bit[x];
	return ret;
}

ll get(int l, int r){
	return get(r) - get(l - 1);
}

int getKth(int x, int k){
	for(int i = 0; i < logN; i++) if(k >> i & 1) x = par[i][x];
	return x;
}

int main(){
	int n; sd(n);
	vector<int> V(n);
	segtree<int> stree(V);
	map<int, int> compress; set<int> vals;
	tree T(n);
	for(int i = 0; i < n; i++){
		sd(C[i]);
		vals.insert(C[i]);
	}

	int pos = 0;
	for(int v : vals) compress[v] = pos++;

	for(int i = 0; i < n; i++) C[i] = compress[C[i]];

	for(int i = 1; i < n; i++){
		sd(A[i]); sd(B[i]);
		A[i]--; B[i]--;
		T.add_edge(A[i], B[i]);
		arrival[B[i]] = i;
		par[0][B[i]] = A[i];
		children[A[i]].push_back(B[i]);
	}
	dfs(0, 0);
	T.rootify(0);
	T.hld(0);

	for(int j = 1; j < logN; j++) for(int i = 0; i < n; i++) par[j][i] = par[j - 1][par[j - 1][i]];

	for(int i = 1; i < n; i++){
		int x = B[i];
		int curr = 1;
		vector<pii> vec;
		ll num = 0;
		while(curr <= depth[x]){
			int y = T.kthAncestor(x, curr);
			// int y = getKth(x, curr);
			int latestNode = stree.query(st[y], en[y]); // last added in x's subtree
			// maximum with the same value
			int lo = curr, hi = depth[x];
			while(lo < hi){
				int mid = (lo + hi + 1) >> 1;
				int node = T.kthAncestor(x, mid);
				if(stree.query(st[node], en[node]) == latestNode) lo = mid;
				else hi = mid - 1;
			}
			int v = C[latestNode], cnt = lo - curr + 1;

			vec.push_back({v, cnt}); // v comes cnt times

			num += cnt * (ll) get(0, v - 1); // the values below should be smaller than this

			add(v, cnt);
			curr = lo + 1;
		}
		for(auto it : vec) add(it.F, -it.S); // reinitialize bit to 0
		printf("%lld\n", num);
		
		stree.modify(st[x], x);
	}
}

컴파일 시 표준 에러 (stderr) 메시지

construction.cpp: In member function 'void tree::rootify(int)':
construction.cpp:90:27: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int i = 0; (1 << i) <= tour.size(); i++)
                  ~~~~~~~~~^~~~~~~~~~~~~~
construction.cpp:91:50: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
    for(int j = (1 << i); j < (1 << (i + 1)) && j <= tour.size(); j++)
                                                ~~^~~~~~~~~~~~~~
construction.cpp:101:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int i = 0; i < tour.size(); i += block_size){
                  ~~^~~~~~~~~~~~~
construction.cpp:106:12: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
    for(; j < tour.size() && j < i + block_size; j++){
          ~~^~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:5:20: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
 #define sd(x) scanf("%d", &(x))
               ~~~~~^~~~~~~~~~~~
construction.cpp:328:9: note: in expansion of macro 'sd'
  int n; sd(n);
         ^~
construction.cpp:5:20: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
 #define sd(x) scanf("%d", &(x))
               ~~~~~^~~~~~~~~~~~
construction.cpp:334:3: note: in expansion of macro 'sd'
   sd(C[i]);
   ^~
construction.cpp:5:20: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
 #define sd(x) scanf("%d", &(x))
               ~~~~~^~~~~~~~~~~~
construction.cpp:344:3: note: in expansion of macro 'sd'
   sd(A[i]); sd(B[i]);
   ^~
construction.cpp:5:20: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
 #define sd(x) scanf("%d", &(x))
               ~~~~~^~~~~~~~~~~~
construction.cpp:344:13: note: in expansion of macro 'sd'
   sd(A[i]); sd(B[i]);
             ^~

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