#916749 (C++17) No.2491 Pochi and A Warp Machine

提出ソース

結果

問題	No.2491 Pochi and A Warp Machine
ユーザー	square1001
提出日時	2023-09-29 23:19:35
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	1,697 ms / 3,000 ms
コード長	6,793 bytes
コンパイル時間	4,093 ms
コンパイル使用メモリ	237,532 KB
実行使用メモリ	59,532 KB
最終ジャッジ日時	2024-07-22 17:18:18
合計ジャッジ時間	35,891 ms
ジャッジサーバーID （参考情報）	judge2 / judge5

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 5,248 KB
testcase_01	AC	2 ms 5,376 KB
testcase_02	AC	3 ms 5,376 KB
testcase_03	AC	2 ms 5,376 KB
testcase_04	AC	2 ms 5,376 KB
testcase_05	AC	3 ms 5,376 KB
testcase_06	AC	3 ms 5,376 KB
testcase_07	AC	3 ms 5,376 KB
testcase_08	AC	2 ms 5,376 KB
testcase_09	AC	17 ms 5,376 KB
testcase_10	AC	16 ms 5,376 KB
testcase_11	AC	15 ms 5,376 KB
testcase_12	AC	15 ms 5,376 KB
testcase_13	AC	15 ms 5,376 KB
testcase_14	AC	604 ms 20,996 KB
testcase_15	AC	231 ms 10,832 KB
testcase_16	AC	142 ms 8,412 KB
testcase_17	AC	717 ms 23,648 KB
testcase_18	AC	485 ms 17,996 KB
testcase_19	AC	1,697 ms 45,544 KB
testcase_20	AC	1,681 ms 46,072 KB
testcase_21	AC	1,141 ms 34,040 KB
testcase_22	AC	1,107 ms 33,272 KB
testcase_23	AC	1,101 ms 33,400 KB
testcase_24	AC	1,143 ms 32,704 KB
testcase_25	AC	1,169 ms 35,232 KB
testcase_26	AC	1,098 ms 30,712 KB
testcase_27	AC	1,117 ms 34,680 KB
testcase_28	AC	1,138 ms 33,512 KB
testcase_29	AC	1,110 ms 31,488 KB
testcase_30	AC	1,165 ms 34,556 KB
testcase_31	AC	1,172 ms 34,300 KB
testcase_32	AC	1,126 ms 34,940 KB
testcase_33	AC	1,142 ms 33,404 KB
testcase_34	AC	1,161 ms 34,808 KB
testcase_35	AC	1,149 ms 34,044 KB
testcase_36	AC	1,163 ms 35,192 KB
testcase_37	AC	1,193 ms 34,676 KB
testcase_38	AC	1,102 ms 32,128 KB
testcase_39	AC	1,150 ms 35,192 KB
testcase_40	AC	1,097 ms 31,864 KB
testcase_41	AC	286 ms 59,528 KB
testcase_42	AC	281 ms 59,532 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#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;
	vector<int> comp;
	vector<long long> val;
public:
	fenwicktree() : n(0), comp(vector<int>()), val(vector<long long>()) {}
	fenwicktree(const vector<int>& comp_) : n(comp_.size()), comp(comp_), val(vector<long long>(comp_.size() + 1)) {}
	void add(int pos, long long x) {
		pos = lower_bound(comp.begin(), comp.end(), pos) - comp.begin();
		for (int i = pos + 1; i <= n; i += i & (-i)) {
			val[i] += x;
		}
	}
	long long rangesum(int r) {
		r = lower_bound(comp.begin(), comp.end(), r) - comp.begin();
		long long ans = 0;
		for (int i = r; i >= 1; i -= i & (-i)) {
			ans += val[i];
		}
		return ans;
	}
};

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<int> diff(N);
	for (int i = 0; i < N; i++) {
		int idx = P[i];
		diff[i] = (idx != 0 ? dlist[idx - 1] : 0) - depth[i] - 1;
	}
	int Z = dlist.size() + 1;
	vector<vector<int> > diffcomp(K + 1);
	for (int i = 0; i < N; i++) {
		if (diff[i] >= 0) {
			diffcomp[K].push_back(Z - diff[i] - 1);
			if (i != root) {
				diffcomp[comp[i]].push_back(Z - diff[i] - 1);
			}
		}
	}
	for (int i = 0; i <= K; i++) {
		sort(diffcomp[i].begin(), diffcomp[i].end());
		diffcomp[i].erase(unique(diffcomp[i].begin(), diffcomp[i].end()), diffcomp[i].end());
	}
	vector<fenwicktree> bit0(K + 1);
	vector<fenwicktree> bit1(K + 1);
	for (int i = 0; i <= K; i++) {
		bit0[i] = fenwicktree(diffcomp[i]);
		bit1[i] = fenwicktree(diffcomp[i]);
	}
	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(Z - depth[i]);
		profit -= bit0[K].rangesum(Z - depth[i]) * depth[i];
		if (i != root) {
			profit -= bit1[comp[i]].rangesum(Z - depth[i]);
			profit += bit0[comp[i]].rangesum(Z - depth[i]) * depth[i];
		}
		ans[idx] -= profit;
		int diff = (idx != 0 ? dlist[idx - 1] : 0) - depth[i] - 1;
		// cerr << "i = " << i << ": idx = " << idx << ", profit = " << profit << ", diff = " << diff << endl;
		if (diff >= 0) {
			bit0[K].add(Z - diff - 1, 1);
			bit1[K].add(Z - diff - 1, diff);
			if (i != root) {
				bit0[comp[i]].add(Z - diff - 1, 1);
				bit1[comp[i]].add(Z - diff - 1, 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;
}

yukicoder

結果

テストケース

ソースコード