Submission #5559996

---

Source Code Expand

#pragma GCC optimize ("-O3") 
#include <bits/stdc++.h>
#include<unordered_map>
using namespace std;

using VS = vector<string>;    using LL = long long;
using VI = vector<int>;       using VVI = vector<VI>;
using PII = pair<int, int>;   using PLL = pair<LL, LL>;
using VL = vector<LL>;        using VVL = vector<VL>;

#define ALL(a)  begin((a)),end((a))
#define RALL(a) (a).rbegin(), (a).rend()
#define SZ(a) int((a).size())
#define SORT(c) sort(ALL((c)))
#define RSORT(c) sort(RALL((c)))
#define UNIQ(c) (c).erase(unique(ALL((c))), end((c)))
#define FOR(i, s, e) for (int(i) = (s); (i) < (e); (i)++)
#define FORR(i, s, e) for (int(i) = (s); (i) > (e); (i)--)
#ifdef YANG33
#include "mydebug.hpp"
#else
#define DD(x) 
#endif
const int INF = 1e9;                          const LL LINF = 1e16;
const LL MOD = 1000000007;                    const double PI = acos(-1.0);
int DX[8] = { 0, 0, 1, -1, 1, 1, -1, -1 };    int DY[8] = { 1, -1, 0, 0, 1, -1, 1, -1 };

/* -----  __MAKE_TIME__  Problem: __PROBLEM__ / Link: __CONTEST_URL__  ----- */



template <class N, class E>
struct Rerooting {
	int n;
	const vector<vector<E>>& G;
	vector<N> subtreesum;
	vector<vector<N>> dp; // L
	vector<unordered_map <int, N>>edgememo;
	void dfs1(int v, int p) {
		subtreesum[v] = N();
		for (auto e : G[v]) {
			int nx = e.to;
			if (nx == p) continue;
			dfs1(nx, v);
			subtreesum[v] = subtreesum[v] + subtreesum[nx].include_e(v, e);
		}
		subtreesum[v] = subtreesum[v].include_v(v);
		edgememo[v][p] = subtreesum[v];
	}
	void dfs2(int v, int p, N top) {
		int deg = int(G[v].size());
		dp[v] = vector<N>(deg + 1);
		dp[v][0] = N();
		for (int i = 0; i < deg; i++) {
			int nx = G[v][i].to;
			dp[v][i + 1] =
				dp[v][i] + (nx == p ? top : subtreesum[nx]).include_e(v, G[v][i]);
		}
		N rnode = N();
		dp[v].back() = dp[v].back().include_v(v);
		for (int i = deg - 1; i >= 0; i--) {
			dp[v][i] = (dp[v][i] + rnode).include_v(v);
			int nx = G[v][i].to;
			if (nx != p)edgememo[v][nx] = dp[v][i], dfs2(nx, v, dp[v][i]);
			rnode = rnode + (nx == p ? top : subtreesum[nx]).include_e(v, G[v][i]);
		}
	}
	Rerooting(const vector<vector<E>>& _g, int root) : n(int(_g.size())), G(_g), subtreesum(n), dp(n), edgememo(int(_g.size())) {
		dfs1(root, -1);
		dfs2(root, -1, N());
	}
};

template <class N, class E> vector<vector<N>> Rerooting_get(const vector<vector<E>>& g, int root) {
	return Rerooting<N, E>(g, root).dp;
}
// me[cur][pre]:=:根がpreにあるときの頂点curの部分木の値

inline int ri() { int in; int y = scanf("%d", &in); return in; }
inline LL ril() { LL in; int y = scanf("%lld", &in); return in; }
inline void oi(LL out) { printf("%lld\n", out); }

int main() {

	struct Edge
	{
		int to;
		Edge(int t) :to(t) {}
	};

	int N = ri();
	vector<vector<Edge>>G(N);
	vector<PII>edge(N - 1);
	FOR(i, 0, N - 1) {
		int a = ri() - 1, b = ri() - 1;
		G[a].push_back(Edge(b));
		G[b].push_back(Edge(a));
		edge[i] = PII(a, b);
	}

	static string s; cin >> s;

	struct Node {
		PLL d;
		LL onecnt;
		Node() :d(PLL(0, 0)), onecnt(0) {}
		Node include_e(int, const Edge&) const {
			Node res = *this;
			res.d.second = 0;
			res.d.first++;

			return res;
		}

		Node operator + (const Node& r) const {
			Node res = *this;
			for (auto it : { r.d.first,r.d.second }) {
				if (res.d.first < it) {
					res.d.second = res.d.first;
					res.d.first = it;
				}
				else if (res.d.second < it) {
					res.d.second = it;
				}
			}
			res.onecnt += r.onecnt;
			return res;
		}

		Node include_v(int v) const {
			Node res = *this;
			if (s[v] == '1')res.onecnt++;
			return res;
		}
	};

	auto all_tree = Rerooting<Node, Edge>(G, 0);
	auto tree_dp_v = all_tree.dp;
	auto tree_dp_e = all_tree.edgememo;
	LL ans = 0;

	// vertex c
	FOR(i, 0, N) {
		LL OK = tree_dp_v[i].back().d.second;
		LL DW = INF;
		DD(de(tree_dp_v[i].back().d));
		for (auto it : tree_dp_v[i]) {
			if (it.onecnt)
				DW = min(DW, it.d.first);
		}
		if (s[i] == '1')DW = 0;
		DD(de(OK, DW));
		if (DW <= OK)
			ans += OK - DW + 1;
	}
	DD(de(ans));
	// edge c
	FOR(i, 0, N - 1) {
		int a, b;
		tie(a, b) = edge[i];
		Node d1 = tree_dp_e[a][b];
		Node d2 = tree_dp_e[b][a];
		DD(de(d1.d, d1.onecnt, d2.d, d2.onecnt));
		if (d1.d.first > d2.d.first) {
			ans += !!(d2.onecnt > 0);
		}
		else if (d1.d.first < d2.d.first) {
			ans += !!(d1.onecnt > 0);
		}
		else ans++;

	}
	oi(ans);

	//FOR(i, 0, N) {
	//	for (auto it : tree_dp_e[i]) {
	//		cout << i << "," << it.first << "," << it.second.d.first << "," << (it.second.onecnt>0) << endl;
	//	}
	//}
	return 0;
}

Submission Info

Judge Result