#1160260 (C++23) No.2892 Lime and Karin

• 問題: No.2892 Lime and Karin
• ユーザー: kwm_t
• 提出日時: 2026-04-19 00:58:55
• 言語: C++23 (gcc 15.2.0 + boost 1.89.0)
• 結果: AC
• 実行時間: 118 ms / 8,000 ms
• コード長: 5,310 bytes
• コンパイル時間: 6,129 ms
• コンパイル使用メモリ: 382,408 KB
• 実行使用メモリ: 29,056 KB
• 最終ジャッジ日時: 2026-04-19 00:59:40
• 合計ジャッジ時間: 16,595 ms
• ジャッジサーバーID （参考情報）: judge1_1 / judge2_1

ファイルパターン	結果
sample	AC * 3
other	AC * 52

ソースコード

#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; }
#ifndef KWM_T_ALGORITHM_DSU_ON_TREE_HPP
#define KWM_T_ALGORITHM_DSU_ON_TREE_HPP

#include <vector>
#include <utility>

/**
 * @brief DSU on Tree (Sack Technique)
 *
 * @details
 *  木上で部分木クエリを高速に処理するテクニック。
 *  Heavy Child を残し、Small Child を都度破棄することで
 *  計算量を削減する。
 *
 *  Euler Tour により部分木を区間 [from[u], to[u]) に対応させる。
 *
 *  提供機能:
 *    - run(): 各頂点の部分木に対するクエリ
 *    - run_every_pair(): 全頂点ペアに関するクエリ
 *
 *  計算量:
 *    add/erase:O(NlogN)回
 *    query/reset:N回
 *
 *  verified:
 */
namespace kwm_t::algorithm {

class DsuOnTree {
private:
	std::vector<std::vector<int>> graph;
	int n;
	int root;

	std::vector<int> sub;
	std::vector<int> ord;
	std::vector<int> from;
	std::vector<int> to;

	int dfs_sub(int v, int p) {
		sub[v] = 1;

		if (graph[v].size() >= 2 && graph[v][0] == p) {
			std::swap(graph[v][0], graph[v][1]);
		}

		for (int i = 0; i < (int)graph[v].size(); i++) {
			int u = graph[v][i];
			if (u == p) continue;

			sub[v] += dfs_sub(u, v);

			if (i != 0 && sub[graph[v][0]] < sub[u]) {
				std::swap(graph[v][0], graph[v][i]);
			}
		}

		return sub[v];
	}

	void dfs_ord(int v, int p, int& idx) {
		from[v] = idx;
		ord[idx++] = v;

		for (int u : graph[v]) {
			if (u == p) continue;
			dfs_ord(u, v, idx);
		}

		to[v] = idx;
	}

public:
	DsuOnTree(const std::vector<std::vector<int>>& g, int r)
		: graph(g), n((int)g.size()), root(r),
		sub(n), ord(n), from(n), to(n) {
		dfs_sub(root, -1);
		int idx = 0;
		dfs_ord(root, -1, idx);
	}

	template <typename Update, typename Query, typename Clear, typename Reset>
	void run(Update& update, Query& query, Clear& clear, Reset& reset) {
		auto dfs = [&](auto&& self, int v, int p, bool keep) -> void {
			for (int i = 1; i < (int)graph[v].size(); i++) {
				int u = graph[v][i];
				if (u == p) continue;
				self(self, u, v, false);
			}

			if (sub[v] > 1) {
				self(self, graph[v][0], v, true);
			}

			if (sub[v] > 1) {
				int heavy = graph[v][0];
				for (int i = to[heavy]; i < to[v]; i++) {
					update(ord[i]);
				}
			}

			update(v);
			query(v);

			if (!keep) {
				for (int i = from[v]; i < to[v]; i++) {
					clear(ord[i]);
				}
				reset();
			}
		};

		dfs(dfs, root, -1, false);
	}

	template <typename Update, typename Query_Root, typename Query_Light, typename Clear, typename Reset>
	void run_every_pair(Update& update, Query_Root& query_root, Query_Light& query_light, Clear& clear, Reset& reset) {

		auto dfs = [&](auto&& self, int v, int p, bool keep) -> void {
			for (int i = 1; i < (int)graph[v].size(); i++) {
				int u = graph[v][i];
				if (u == p) continue;
				self(self, u, v, false);
			}

			if (sub[v] > 1) {
				self(self, graph[v][0], v, true);
			}

			if (sub[v] > 1) {
				for (int i = 1; i < (int)graph[v].size(); i++) {
					int ch = graph[v][i];
					if (ch == p) continue;

					for (int j = from[ch]; j < to[ch]; j++) {
						query_light(v, ord[j]);
					}
					for (int j = from[ch]; j < to[ch]; j++) {
						update(ord[j]);
					}
				}
			}

			update(v);
			query_root(v);

			if (!keep) {
				for (int i = from[v]; i < to[v]; i++) {
					clear(ord[i]);
				}
				reset();
			}
		};

		dfs(dfs, root, -1, false);
	}
};

} // namespace kwm_t::algorithm

#endif // KWM_T_ALGORITHM_DSU_ON_TREE_HPP

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);
	};
	auto reset = [&]() {
		// nop
	};
	kwm_t::algorithm::DsuOnTree dot(g, 0);
	dot.run_every_pair(update, query_root, query_light, clear, reset);

	cout << ans << endl;
	return 0;
}