結果

問題 No.2491 Pochi and A Warp Machine
コンテスト
ユーザー square1001
提出日時 2023-09-29 22:53:43
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
WA  
実行時間 -
コード長 5,985 bytes
コンパイル時間 3,086 ms
コンパイル使用メモリ 232,504 KB
最終ジャッジ日時 2025-02-17 03:29:05
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2 WA * 1
other AC * 1 WA * 15 TLE * 1 -- * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

const int INF = 1012345678;

struct edge {
	int va, vb;
};

struct state {
	int depth, current, bottom;
};

string to_string(const vector<int>& arr) {
	string res = "[";
	for (int i = 0; i < arr.size(); i++) {
		if (i != 0) {
			res += ", ";
		}
		res += to_string(arr[i]);
	}
	res += "]";
	return res;
}

class fenwicktree {
private:
	int n;
	map<int, long long> d;
	long long rangesum(int r) {
		long long ans = 0;
		for (int i = r; i >= 1; i -= i & (-i)) {
			ans += d[i];
		}
		return ans;
	}
public:
	fenwicktree() : n(0), d(map<int, long long>()) {}
	fenwicktree(int n_) : n(n_), d(map<int, long long>()) {}
	void add(int pos, long long x) {
		for (int i = pos + 1; i <= n; i += i & (-i)) {
			d[i] += x;
		}
	}
	long long rangesum(int l, int r) {
		if (l >= r) {
			return 0;
		}
		return rangesum(r) - rangesum(l);
	}
};

void solve(int N, const vector<int>& P, const vector<edge>& E, const vector<int>& dlist, vector<long long>& ans) {
	// step #1. make graph
	vector<vector<int> > G(N);
	for (edge e : E) {
		G[e.va].push_back(e.vb);
		G[e.vb].push_back(e.va);
	}

	// step #2. find centroid
	int root = -1;
	vector<int> subsize(N, 1);
	auto find_centroid = [&](auto& self, int pos, int pre) -> void {
		bool flag = true;
		for (int i : G[pos]) {
			if (i != pre) {
				self(self, i, pos);
				subsize[pos] += subsize[i];
				if (subsize[i] * 2 > N) {
					flag = false;
				}
			}
		}
		if (flag && subsize[pos] * 2 >= N) {
			root = pos;
		}
	};
	find_centroid(find_centroid, 0, -1);

	// step #3. build tree
	vector<int> depth(N);
	vector<int> comp(N, -1);
	vector<vector<int> > child(N);
	auto build_tree = [&](auto& self, int pos, int pre) -> void {
		for (int i : G[pos]) {
			if (i != pre) {
				comp[i] = comp[pos];
				depth[i] = depth[pos] + 1;
				child[pos].push_back(i);
				self(self, i, pos);
			}
		}
	};
	int K = G[root].size();
	child[root] = G[root];
	for (int i = 0; i < K; i++) {
		comp[G[root][i]] = i;
		depth[G[root][i]] = 1;
		build_tree(build_tree, G[root][i], root);
	}

	// step #5. calculate answer passing root
	vector<int> ord(N);
	for (int i = 0; i < N; i++) {
		ord[i] = i;
	}
	sort(ord.begin(), ord.end(), [&](int va, int vb) {
		return P[va] > P[vb];
	});
	// cerr << "N = " << N << endl;
	// cerr << "centroid = " << root << endl;
	vector<fenwicktree> bit0(K + 1, fenwicktree(N));
	vector<fenwicktree> bit1(K + 1, fenwicktree(N));
	for (int i : ord) {
		int idx = P[i];
		// cost1: dlist[target_idx - 1]
		// cost2: 1 + depth[i] + depth[target]
		// if dlist[target_idx - 1] - depth[target] - 1 >= depth[i], profit
		long long profit = 0;
		profit += bit1[K].rangesum(depth[i], N);
		profit -= bit0[K].rangesum(depth[i], N) * depth[i];
		if (i != root) {
			profit -= bit1[comp[i]].rangesum(depth[i], N);
			profit += bit0[comp[i]].rangesum(depth[i], N) * depth[i];
		}
		ans[idx] -= profit;
		int diff = (idx != 0 ? dlist[idx - 1] : 0) - depth[i] - 1;
		// cerr << "i = " << i << ": profit = " << profit << ", diff = " << diff << endl;
		if (diff >= 0) {
			bit0[K].add(diff, 1);
			bit1[K].add(diff, diff);
			if (i != root) {
				bit0[comp[i]].add(diff, 1);
				bit1[comp[i]].add(diff, diff);
			}
		}
	}

	// step #6. divide into subproblems
	vector<vector<int> > subcomp(K);
	for (int i = 0; i < N; i++) {
		if (i != root) {
			subcomp[comp[i]].push_back(i);
		}
	}
	vector<vector<int> > subp(K);
	vector<vector<edge> > sube(K);
	for (int i = 0; i < K; i++) {
		subp[i].resize(subcomp[i].size());
		for (int j = 0; j < subcomp[i].size(); j++) {
			subp[i][j] = P[subcomp[i][j]];
		}
	}
	for (int i = 0; i < N - 1; i++) {
		if (comp[E[i].va] == comp[E[i].vb]) {
			int id = comp[E[i].va];
			int nva = lower_bound(subcomp[id].begin(), subcomp[id].end(), E[i].va) - subcomp[id].begin();
			int nvb = lower_bound(subcomp[id].begin(), subcomp[id].end(), E[i].vb) - subcomp[id].begin();
			sube[id].push_back(edge{nva, nvb});
		}
	}
	for (int i = 0; i < K; i++) {
		solve(subcomp[i].size(), subp[i], sube[i], dlist, ans);
	}
}

vector<int> distances(int N, const vector<edge>& E, const vector<int>& S, const vector<int>& T) {
	vector<vector<int> > G(N);
	for (edge e : E) {
		G[e.va].push_back(e.vb);
		G[e.vb].push_back(e.va);
	}
	vector<int> par(N);
	vector<int> depth(N);
	auto build_tree = [&](auto& self, int pos, int pre) -> void {
		for (int i : G[pos]) {
			if (i != pre) {
				par[i] = pos;
				depth[i] = depth[pos] + 1;
				self(self, i, pos);
			}
		}
	};
	build_tree(build_tree, 0, -1);
	int bits = (N >= 3 ? 32 - __builtin_clz(N - 1) : N);
	vector<vector<int> > spar(bits);
	spar[0] = par;
	spar[0][0] = 0;
	for (int i = 1; i < bits; i++) {
		spar[i].resize(N);
		for (int j = 0; j < N; j++) {
			spar[i][j] = spar[i - 1][spar[i - 1][j]];
		}
	}
	auto lca = [&](int va, int vb) {
		if (depth[va] < depth[vb]) {
			swap(va, vb);
		}
		int d = depth[va] - depth[vb];
		for (int i = 0; i < bits; i++) {
			if ((d >> i) & 1) {
				va = spar[i][va];
			}
		}
		if (va == vb) {
			return va;
		}
		for (int i = bits - 1; i >= 0; i--) {
			if (spar[i][va] != spar[i][vb]) {
				va = spar[i][va];
				vb = spar[i][vb];
			}
		}
		return spar[0][va];
	};
	vector<int> ans(S.size());
	for (int i = 0; i < S.size(); i++) {
		ans[i] = depth[S[i]] + depth[T[i]] - depth[lca(S[i], T[i])] * 2;
		// cerr << "N = " << N << ", S = " << S[i] << ", T = " << T[i] << ", ans = " << ans[i] << endl;
	}
	return ans;
}

int main() {
	cin.tie(0);
	ios::sync_with_stdio(false);
	int N;
	cin >> N;
	vector<edge> E(N - 1);
	for (int i = 0; i < N - 1; i++) {
		cin >> E[i].va >> E[i].vb;
		E[i].va -= 1;
		E[i].vb -= 1;
	}
	vector<int> P(N), A(N - 1), B(N - 1);
	for (int i = 0; i < N; i++) {
		P[i] = i;
		if (i != N - 1) {
			A[i] = i;
			B[i] = i + 1;
		}
	}
	vector<int> D = distances(N, E, A, B);
	long long dsum = 0;
	for (int i = 0; i < N - 1; i++) {
		dsum += D[i];
	}
	vector<long long> ans(N, dsum);
	solve(N, P, E, D, ans);
	for (int i = 0; i < N; i++) {
		cout << ans[i] << '\n';
	}
	return 0;
}
0