結果

問題 No.2892 Lime and Karin
ユーザー kwm_t
提出日時 2025-02-01 16:42:30
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 175 ms / 8,000 ms
コード長 4,306 bytes
コンパイル時間 10,435 ms
コンパイル使用メモリ 333,744 KB
実行使用メモリ 28,160 KB
最終ジャッジ日時 2025-02-01 16:42:49
合計ジャッジ時間 16,313 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 52
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
//using mint = modint1000000007;
//const int mod = 1000000007;
//using mint = modint998244353;
//const int mod = 998244353;
//const int INF = 1e9;
//const long long LINF = 1e18;
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define rep2(i,l,r)for(int i=(l);i<(r);++i)
#define rrep(i, n) for (int i = (n) - 1; i >= 0; --i)
#define rrep2(i,l,r)for(int i=(r) - 1;i>=(l);--i)
#define all(x) (x).begin(),(x).end()
#define allR(x) (x).rbegin(),(x).rend()
#define P pair<int,int>
template<typename A, typename B> inline bool chmax(A & a, const B & b) { if (a < b) { a = b; return true; } return false; }
template<typename A, typename B> inline bool chmin(A & a, const B & b) { if (a > b) { a = b; return true; } return false; }
template<typename G>
class DsuOnTree {
public:
	DsuOnTree(G &g, int root = 0)
		:g(g),
		n(g.size()),
		sub_size(g.size()),
		euler(g.size()),
		down(g.size()),
		up(g.size()),
		idx_(0),
		root(root) {
		dfs1(root);
		dfs2(root);
	}

	template <typename UPDATE, typename QUERY, typename CLEAR, typename RESET>
	void run(UPDATE &update, QUERY &query, CLEAR &clear, RESET &reset) {
		auto dsu = [&](auto &&self, int v, int p, bool keep)->void {
			// un heavy path
			for (int i = 1; i < (int)g[v].size(); ++i) {
				if (g[v][i] == p) continue;
				self(self, g[v][i], v, false);
			}
			// not leaf
			if (sub_size[v] != 1) {
				self(self, g[v][0], v, true);
				for (int i = up[g[v][0]]; i < up[v]; i++) {
					update(v, euler[i]);
				}
			}
			update(v, v);
			query(v);
			if (!keep) {
				for (int i = down[v]; i < up[v]; i++) clear(euler[i]);
				reset(v);
			}
		};
		dsu(dsu, root, -1, true);
	}

	using F = function<void(int)>;
	using F2 = function<void(int, int)>;
	void run_every_pair(F update, F query_root, F2 query_light, F clear) {
		auto dfs = [&](auto &&self, int v, int p, bool keep) -> void {
			// un heavy path
			for (int i = 1; i < (int)g[v].size(); i++) {
				if (g[v][i] == p) continue;
				self(self, g[v][i], v, false);
			}

			// not leaf
			if (sub_size[v] != 1) {
				self(self, g[v][0], v, true);
			}

			// update Small Child
			if (sub_size[v] != 1) {
				for (int i = 1; i < g[v].size(); i++) {
					if (g[v][i] == p) continue;

					int ch = g[v][i];
					// Big Child とすでに処理済みの Small Child に関するクエリを処理
					for (int u = down[ch]; u < up[ch]; u++) query_light(v, euler[u]);
					for (int u = down[ch]; u < up[ch]; u++) update(euler[u]);
				}
			}

			update(v);
			query_root(v);

			if (!keep) {
				for (int i = down[v]; i < up[v]; i++) clear(euler[i]);
			}
			return;
		};

		dfs(dfs, root, -1, true);
	}
private:
	G &g;
	int n;
	std::vector<int>sub_size, euler, down, up;
	int idx_;
	int root;

	int dfs1(int v, int p = -1) {
		sub_size[v] = 1;
		// 先頭にheavy pathを
		if ((int)g[v].size() >= 2 && g[v][0] == p) {
			std::swap(g[v][0], g[v][1]);
		}
		for (auto &e : g[v]) {
			if (e == p)continue;
			sub_size[v] += dfs1(e, v);
			if (sub_size[e] > sub_size[g[v][0]]) {
				std::swap(g[v][0], e);
			}
		}
		return sub_size[v];
	}
	void dfs2(int v, int p = -1) {
		euler[idx_] = v;
		down[v] = idx_++;
		for (auto &e : g[v]) {
			if (e == p)continue;
			dfs2(e, v);
		}
		up[v] = idx_;
	}
};
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	int n; cin >> n;
	vector<vector<int>>g(n);
	rep(I, n - 1) {
		int u, v; cin >> u >> v;
		u--, v--;
		g[u].push_back(v);
		g[v].push_back(u);
	}
	string s;
	cin >> s;
	vector<int> a(n, -1);
	rep(i, n) if (s[i] == '1')a[i] = 1;

	vector<int> score(n);
	auto dfs = [&](auto &&f, int v, int p, int d) -> void {
		score[v] = d + a[v];
		for (int to : g[v]) {
			if (to == p) continue;
			f(f, to, v, score[v]);
		}
	};
	dfs(dfs, 0, -1, 0);

	fenwick_tree<long long> fw(2 * n + 1);
	long long ans = 0;
	auto update = [&](int v) {
		fw.add(score[v] + n, 1);
	};
	auto query_root = [&](int v) {
		int p = score[v] - a[v];
		ans += fw.sum(p + 1 + n, 2 * n + 1);
	};
	auto query_light = [&](int r, int v) {
		ans += fw.sum(2 * score[r] - a[r] - score[v] + 1 + n, 2 * n + 1);
	};
	auto clear = [&](int v) {
		fw.add(score[v] + n, -1);
	};
	DsuOnTree dsu(g, 0);
	dsu.run_every_pair(update, query_root, query_light, clear);
	cout << ans << endl;
	return 0;
}
0