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

提出ソース

結果

問題	No.2491 Pochi and A Warp Machine
ユーザー	square1001
提出日時	2023-09-29 23:07:39
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	TLE
実行時間	-
コード長	6,868 bytes
コンパイル時間	3,005 ms
コンパイル使用メモリ	242,660 KB
実行使用メモリ	98,484 KB
最終ジャッジ日時	2023-09-29 23:08:45
合計ジャッジ時間	52,653 ms
ジャッジサーバーID （参考情報）	judge11 / judge13

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 4,376 KB
testcase_01	AC	1 ms 4,380 KB
testcase_02	AC	1 ms 4,376 KB
testcase_03	AC	1 ms 4,376 KB
testcase_04	AC	2 ms 4,380 KB
testcase_05	AC	2 ms 4,380 KB
testcase_06	AC	1 ms 4,376 KB
testcase_07	AC	2 ms 4,384 KB
testcase_08	AC	2 ms 4,376 KB
testcase_09	AC	21 ms 4,376 KB
testcase_10	AC	21 ms 4,380 KB
testcase_11	AC	21 ms 4,376 KB
testcase_12	AC	22 ms 4,380 KB
testcase_13	AC	21 ms 4,380 KB
testcase_14	AC	947 ms 21,280 KB
testcase_15	AC	365 ms 11,028 KB
testcase_16	AC	224 ms 8,752 KB
testcase_17	AC	1,146 ms 24,096 KB
testcase_18	AC	810 ms 18,640 KB
testcase_19	TLE	-
testcase_20	TLE	-
testcase_21	AC	1,714 ms 34,300 KB
testcase_22	AC	1,736 ms 33,616 KB
testcase_23	AC	1,749 ms 33,868 KB
testcase_24	AC	1,768 ms 33,308 KB
testcase_25	AC	1,847 ms 35,932 KB
testcase_26	AC	1,805 ms 31,460 KB
testcase_27	AC	1,772 ms 35,452 KB
testcase_28	AC	1,794 ms 34,000 KB
testcase_29	AC	1,773 ms 32,136 KB
testcase_30	AC	1,772 ms 35,208 KB
testcase_31	AC	1,832 ms 34,788 KB
testcase_32	AC	1,749 ms 35,092 KB
testcase_33	AC	1,800 ms 33,916 KB
testcase_34	AC	1,784 ms 35,100 KB
testcase_35	AC	1,795 ms 34,676 KB
testcase_36	AC	1,812 ms 35,952 KB
testcase_37	AC	1,872 ms 35,396 KB
testcase_38	AC	1,731 ms 32,428 KB
testcase_39	AC	1,823 ms 35,664 KB
testcase_40	AC	1,692 ms 32,500 KB
testcase_41	AC	437 ms 98,484 KB
testcase_42	AC	408 ms 98,224 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;
	unordered_map<int, long long> d;
	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 += d[i];
		}
		return ans;
	}
public:
	fenwicktree() : n(0), d(unordered_map<int, 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)) {
			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<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;
	}
	vector<vector<int> > diffcomp(K + 1);
	for (int i = 0; i < N; i++) {
		if (diff[i] >= 0) {
			diffcomp[K].push_back(diff[i]);
			if (i != root) {
				diffcomp[comp[i]].push_back(diff[i]);
			}
		}
	}
	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());
	}
	int Z = dlist.size() + 1;
	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(depth[i], Z);
		profit -= bit0[K].rangesum(depth[i], Z) * depth[i];
		if (i != root) {
			profit -= bit1[comp[i]].rangesum(depth[i], Z);
			profit += bit0[comp[i]].rangesum(depth[i], Z) * 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(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;
}

yukicoder

結果

テストケース

ソースコード