結果
問題 | No.2892 Lime and Karin |
ユーザー | tonegawa |
提出日時 | 2024-09-15 13:21:26 |
言語 | C++17(gcc12) (gcc 12.3.0 + boost 1.87.0) |
結果 |
AC
|
実行時間 | 1,107 ms / 8,000 ms |
コード長 | 21,155 bytes |
コンパイル時間 | 2,898 ms |
コンパイル使用メモリ | 177,284 KB |
実行使用メモリ | 105,684 KB |
最終ジャッジ日時 | 2024-09-15 13:21:56 |
合計ジャッジ時間 | 29,246 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge2 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 52 |
ソースコード
#include <string>#include <iostream>#include <vector>#include <array>#include <tuple>#include <stack>#include <queue>#include <deque>#include <algorithm>#include <set>#include <map>#include <unordered_set>#include <unordered_map>#include <bitset>#include <cmath>#include <functional>#include <cassert>#include <climits>#include <iomanip>#include <numeric>#include <memory>#include <random>#include <thread>#include <chrono>#define allof(obj) (obj).begin(), (obj).end()#define range(i, l, r) for(int i=l;i<r;i++)#define unique_elem(obj) obj.erase(std::unique(allof(obj)), obj.end())#define bit_subset(i, S) for(int i=S, zero_cnt=0;(zero_cnt+=i==S)<2;i=(i-1)&S)#define bit_kpop(i, n, k) for(int i=(1<<k)-1,x_bit,y_bit;i<(1<<n);x_bit=(i&-i),y_bit=i+x_bit,i=(!i?(1<<n):((i&~y_bit)/x_bit>>1)|y_bit))#define bit_kth(i, k) ((i >> k)&1)#define bit_highest(i) (i?63-__builtin_clzll(i):-1)#define bit_lowest(i) (i?__builtin_ctzll(i):-1)#define sleepms(t) std::this_thread::sleep_for(std::chrono::milliseconds(t))using ll = long long;using ld = long double;using ul = uint64_t;using pi = std::pair<int, int>;using pl = std::pair<ll, ll>;using namespace std;template<typename F, typename S>std::ostream &operator << (std::ostream &dest, const std::pair<F, S> &p) {dest << p.first << ' ' << p.second;return dest;}template<typename A, typename B>std::ostream &operator << (std::ostream &dest, const std::tuple<A, B> &t) {dest << std::get<0>(t) << ' ' << std::get<1>(t);return dest;}template<typename A, typename B, typename C>std::ostream &operator << (std::ostream &dest, const std::tuple<A, B, C> &t) {dest << std::get<0>(t) << ' ' << std::get<1>(t) << ' ' << std::get<2>(t);return dest;}template<typename A, typename B, typename C, typename D>std::ostream &operator << (std::ostream &dest, const std::tuple<A, B, C, D> &t) {dest << std::get<0>(t) << ' ' << std::get<1>(t) << ' ' << std::get<2>(t) << ' ' << std::get<3>(t);return dest;}template<typename T>std::ostream &operator << (std::ostream &dest, const std::vector<std::vector<T>> &v) {int sz = v.size();if (!sz) return dest;for (int i = 0; i < sz; i++) {int m = v[i].size();for (int j = 0; j < m; j++) dest << v[i][j] << (i != sz - 1 && j == m - 1 ? '\n' : ' ');}return dest;}template<typename T>std::ostream &operator << (std::ostream &dest, const std::vector<T> &v) {int sz = v.size();if (!sz) return dest;for (int i = 0; i < sz - 1; i++) dest << v[i] << ' ';dest << v[sz - 1];return dest;}template<typename T, size_t sz>std::ostream &operator << (std::ostream &dest, const std::array<T, sz> &v) {if (!sz) return dest;for (int i = 0; i < sz - 1; i++) dest << v[i] << ' ';dest << v[sz - 1];return dest;}template<typename T>std::ostream &operator << (std::ostream &dest, const std::set<T> &v) {for (auto itr = v.begin(); itr != v.end();) {dest << *itr;itr++;if (itr != v.end()) dest << ' ';}return dest;}template<typename T, typename E>std::ostream &operator << (std::ostream &dest, const std::map<T, E> &v) {for (auto itr = v.begin(); itr != v.end(); ) {dest << '(' << itr->first << ", " << itr->second << ')';itr++;if (itr != v.end()) dest << '\n';}return dest;}template<typename T>vector<T> make_vec(size_t sz, T val) { return std::vector<T>(sz, val); }template<typename T, typename... Tail>auto make_vec(size_t sz, Tail ...tail) {return std::vector<decltype(make_vec<T>(tail...))>(sz, make_vec<T>(tail...));}template<typename T>vector<T> read_vec(size_t sz) {std::vector<T> v(sz);for (int i = 0; i < (int)sz; i++) std::cin >> v[i];return v;}template<typename T, typename... Tail>auto read_vec(size_t sz, Tail ...tail) {auto v = std::vector<decltype(read_vec<T>(tail...))>(sz);for (int i = 0; i < (int)sz; i++) v[i] = read_vec<T>(tail...);return v;}// x / y以上の最小の整数ll ceil_div(ll x, ll y) {assert(y > 0);return (x + (x > 0 ? y - 1 : 0)) / y;}// x / y以下の最大の整数ll floor_div(ll x, ll y) {assert(y > 0);return (x + (x > 0 ? 0 : -y + 1)) / y;}void io_init() {std::cin.tie(nullptr);std::ios::sync_with_stdio(false);}template<typename _W>struct edge_base {using W = _W;int t;edge_base() {}edge_base(int _t) : t(_t) {}virtual int from() const {assert(false);}int to() const {return t;}virtual W weight() const {assert(false);}operator int() const {return t;}};struct simple_edge : edge_base<int> {simple_edge() {}simple_edge(int _t) : edge_base<int>(_t) {}int weight() const override {return 1;}};template<typename _W>struct weighted_edge : edge_base<_W> {using W = _W;W w;weighted_edge() {}weighted_edge(int _t, _W _w) : edge_base<_W>(_t), w(_w) {}W weight() const override {return w;}};template <class E>struct csr {std::vector<int> start;std::vector<E> elist;csr() {}explicit csr(int n, const std::vector<std::pair<int, E>>& edges) : start(n + 1), elist(edges.size()) {for (auto e : edges) {start[e.first + 1]++;}for (int i = 1; i <= n; i++) {start[i] += start[i - 1];}auto counter = start;for (auto e : edges) {elist[counter[e.first]++] = e.second;}}~csr() {}int N() const {return start.size() - 1;}int deg(int i) const {return start[i + 1] - start[i];}int begin(int i) const {return start[i];}int end(int i) const {return start[i + 1];}E& operator [] (int i) { return elist[i]; }};template<typename edge = int>struct hld {int N, root;csr<edge> G;std::vector<int> sz, dep, par, head;std::vector<int> tour, in, out;hld() : N(0) {}hld(int root, const csr<edge> &_G) : N(_G.N()), root(root), G(_G), sz(N), dep(N), par(N), head(N), tour(N), in(N), out(N) {auto dfs_sz = [&](auto &&dfs_sz, int v, int p, int d) -> void {sz[v] = 1, dep[v] = d, par[v] = p;int L = G.begin(v);for (int i = L; i < G.end(v); i++) {int to = G.elist[i];if (to == p) continue;dfs_sz(dfs_sz, to, v, d + 1);sz[v] += sz[to];int hc = G.elist[L];if (hc == p || sz[hc] < sz[to]) {std::swap(G.elist[L], G.elist[i]);}}};int t = 0;auto dfs_tour = [&](auto &&dfs_tour, int v, int p) -> void {in[v] = t;tour[t++] = v;int L = G.begin(v);for (int i = L; i < G.end(v); i++) {int to = G.elist[i];if (to == p) continue;if (i == L) {head[to] = head[v];} else {head[to] = to;}dfs_tour(dfs_tour, to, v);}out[v] = t;};head[root] = root;dfs_sz(dfs_sz, root, -1, 0);dfs_tour(dfs_tour, root, -1);}// k個上の祖先, k > depth[v]なら-1int la(int v, int k) {if (dep[v] < k) return -1;while (true) {int u = head[v];if (in[v] - in[u] >= k) return tour[in[v] - k];k -= in[v] - in[u] + 1;v = par[u];}}int lca(int u, int v) {while (true) {if (in[u] > in[v]) std::swap(u, v);if (head[u] == head[v]) return u;v = par[head[v]];}}int dist(int s, int t) {int p = lca(s, t);return dep[s] + dep[t] - 2 * dep[p];}// 0 <= k <= dist(s, t) でない場合-1int jump_on_tree(int s, int t, int k) {int p = lca(s, t);int sp = dep[s] - dep[p];if (k <= sp) {return la(s, k);} else {int pt = dep[t] - dep[p];if (k > sp + pt) return -1;return la(t, sp + pt - k);}}// pの部分木がcを含むかbool contain_subtree(int p, int c) {if (dep[p] > dep[c]) return false;return p == la(c, dep[c] - dep[p]);}// s->tパス(端点も含む)がvを含むかbool contain_path(int s, int t, int v){if (lca(s, t) == v) return true;return contain_subtree(v, s) ^ contain_subtree(v, t);}int ord_point(int v) {return in[v];}// [l, r)std::pair<int, int> ord_subtree(int v) {return {in[v], out[v]};}// O(logN)個の区間でパスを表す// 各区間はl < rまたはl > rを満たし// l < rなら上から下に降りる順番// l > rなら下から上に登る順番// (s -> t パスの順序を保っている)std::vector<std::pair<int, int>> ord_path(int s, int t) {std::vector<std::pair<int, int>> up, down;while (head[s] != head[t]) {if (in[s] < in[t]) {down.push_back({in[head[t]], in[t] + 1});t = par[head[t]];} else {up.push_back({in[s] + 1, in[head[s]]});s = par[head[s]];}}if (in[s] < in[t]) {down.push_back({in[s], in[t] + 1});} else {up.push_back({in[s] + 1, in[t]});}std::reverse(down.begin(), down.end());up.insert(up.end(), down.begin(), down.end());return up;}// l < rなら上から下に降りる順番// l > rなら下から上に登る順番// (s -> t パスの順序を保っていない, 区間と向きは合っている)template<typename F>void query_path(int s, int t, F f) {while (head[s] != head[t]) {if (in[s] < in[t]) {f(in[head[t]], in[t] + 1);t = par[head[t]];} else {f(in[s] + 1, in[head[s]]);s = par[head[s]];}}if (in[s] < in[t]) {f(in[s], in[t] + 1);} else {f(in[s] + 1, in[t]);}}// F := [l, r)のモノイド積を得る関数template<typename S, S (*op)(S, S), S (*e)(), S (*flip)(S), typename F>S prod_path(int s, int t, F f) {S sl = e(), sr = e();query_path(s, t, [&](int l, int r) {if (l > r) {sl = op(sl, flip(f(r, l)));} else {sr = op(f(l, r), sr);}});return op(sl, sr);}// path [a,b] と [c,d] の交わり. 空ならば {-1,-1}.std::pair<int, int> intersection_path(int a, int b, int c, int d) {int ab = lca(a, b), ac = lca(a, c), ad = lca(a, d);int bc = lca(b, c), bd = lca(b, d), cd = lca(c, d);int x = ab ^ ac ^ bc, y = ab ^ ad ^ bd;if (x != y) return {x, y};int z = ac ^ ad ^ cd;if (x != z) x = -1;return {x, x};}// {頂点集合(in順), 辺集合}std::pair<std::vector<int>, std::vector<std::pair<int, int>>> auxiliary_tree(std::vector<int> v) {if (v.empty()) return {{}, {}};std::sort(v.begin(), v.end(), [&](int a, int b) {return in[a] < in[b]; });std::vector<int> V;std::vector<std::pair<int, int>> E;std::vector<int> st;st.push_back(v[0]);for (int i = 1; i < v.size(); i++) {if (v[i] == v[i - 1]) continue;int l = lca(v[i], v[i - 1]);while (true) {int c = st.back();st.pop_back();if (st.empty() || dep[st.back()] < dep[l]) {st.push_back(l);st.push_back(v[i]);if (dep[c] > dep[l]) {E.push_back({l, c});V.push_back(c);}break;}E.push_back({st.back(), c});V.push_back(c);}}while (st.size() >= 2) {int c = st.back();st.pop_back();int p = st.back();E.push_back({p, c});V.push_back(c);}V.push_back(st[0]);return {V, E};}};// 重心を再帰的に繋いだ木を作るtemplate<typename edge = int>struct centroid {int N;csr<int> G;std::vector<int> sz, dep, par;centroid() : N(0) {}centroid(const csr<edge> &g) : N(g.N()), sz(N, 0), dep(N, N), par(N, -1) {std::vector<std::pair<int, int>> E;auto add_edge = [&](int p, int c) -> void {E.push_back({p, c});par[c] = p;};auto dfs = [&](auto &&dfs, int v, int p, const int M, const int rank) -> std::pair<int, int> {int size = 1;for (int i = g.begin(v); i < g.end(v); i++) {int to = g.elist[i];if (to == p || dep[to] < rank) continue;auto [sz_c, cent_c] = dfs(dfs, to, v, M, rank);if (sz_c == -1) return {-1, cent_c};sz[to] = sz_c;size += sz_c;}if(size * 2 >= M){sz[v] = M;dep[v] = rank;for (int i = g.begin(v); i < g.end(v); i++) {int to = g.elist[i];if (to == p || dep[to] < rank) continue;auto [sz_c, cent_c] = dfs(dfs, to, -1, sz[to], rank + 1);add_edge(v, cent_c);}if (p != -1) {auto [sz_c, cent_c] = dfs(dfs, p, -1, M - size, rank + 1);add_edge(v, cent_c);}return {-1, v};}return {size, -1};};dfs(dfs, 0, -1, N, 0);G = csr<int>(N, E);}// 元の木のパスが重心木でa, bを含む場合lca(a, b)も含む// ランク最小の頂点は1個// s-tパスのランク最小の頂点はlca(s, t)int lca(int a, int b) {if (dep[a] > dep[b]) std::swap(a, b);while (dep[a] < dep[b]) b = par[b];while (a != b) {a = par[a];b = par[b];}return a;}};template<typename Idx, typename T>struct fenwick_tree_sparse {private:int N;std::vector<Idx> I;std::vector<T> V;public:fenwick_tree_sparse() {}// idxの昇順にソート済fenwick_tree_sparse(const std::vector<std::pair<Idx, T>> &v): I(1, std::numeric_limits<Idx>::min()), V(1, 0) {for (int i = 0; i < v.size(); i++) {assert(I.back() <= v[i].first);if (I.back() == v[i].first) {V.back() += v[i].second;} else {I.push_back(v[i].first);V.push_back(v[i].second);}}N = I.size() - 1;for (int i = 1; i <= N; i++) {int nxt = i + (i & (-i));if (nxt <= N) V[nxt] += V[i];}}int size() {return N;}// V[k] <- op(V[k], x)// kを事前に追加していない場合エラーvoid apply(Idx k, T x) {auto itr = std::lower_bound(I.begin(), I.end(), k);assert(itr != I.end() && *itr == k);int l = itr - I.begin();for (int i = l; i <= N; i += (i & (-i))) {V[i] += x;}}// k番目に小さい点にx足すvoid apply_raw(int k, T x) {for (int i = k + 1; i <= N; i += (i & (-i))) {V[i] += x;}}// prod[0, r)T prod(Idx R) {int r = std::lower_bound(I.begin(), I.end(), R) - I.begin() - 1;T res = 0;for (int k = r; k > 0; k -= (k & (-k))) {res = op(V[k], res);}return res;}// 逆元がある必要があるT prod(Idx L, Idx R) {if (L >= R) return 0;int r = std::lower_bound(I.begin(), I.end(), R) - I.begin() - 1;int l = std::lower_bound(I.begin(), I.end(), L) - I.begin() - 1;if (l >= r) return 0;T res = 0;while ((r | l) != l) {res += V[r];r -= (r & (-r));}while (l > r) {res -= V[l];l -= (l & (-l));}return res;}};string S;template<typename T, typename edge>struct contour_sum_weighted {using W = typename edge::W;hld<edge> H;centroid<edge> C;std::vector<int> O;std::vector<std::vector<W>> D;std::vector<fenwick_tree_sparse<W, T>> FT;std::vector<std::vector<fenwick_tree_sparse<W, T>>> FTc;contour_sum_weighted() {}contour_sum_weighted(const hld<edge> &_H, const centroid<edge> &_C, const std::vector<T> &X) : H(_H), C(_C), O(C.N), D(C.N), FT(C.N), FTc(C.N) {int N = C.N;int root = -1;for (int i = 0; i < N; i++) {if (C.par[i] == -1) {root = i;break;}}for (int i = 0; i < N; i++) {D[i].push_back(0);for (int j = C.G.begin(i); j < C.G.end(i); j++) {int k = C.G.elist[j];O[k] = j - C.G.begin(i);}}std::vector<W> wsum(N);wsum[H.root] = (S[H.root] == '1' ? 1 : -1);auto dfs_w = [&](auto &&dfs_w, int v, int p) -> void {for (int i = H.G.begin(v); i < H.G.end(v); i++) {int to = H.G.elist[i];if (to == p) continue;wsum[to] = wsum[v] + H.G.elist[i].weight();dfs_w(dfs_w, to, v);}};dfs_w(dfs_w, H.root, -1);auto dfs = [&](auto &&dfs, int v) -> std::vector<int> {std::vector<int> res{v};std::vector<std::pair<W, T>> tmp{{S[v] == '1' ? 1 : -1, X[v]}};FTc[v].resize(C.G.end(v) - C.G.begin(v));for (int i = C.G.begin(v); i < C.G.end(v); i++) {int c = C.G.elist[i];auto V = dfs(dfs, c);std::vector<std::pair<W, T>> tmpc;for (int c : V) {int L = H.lca(c, v);W dist1 = wsum[c] + wsum[v] - 2 * wsum[L];W dist2 = dist1 + (S[L] == '1' ? 1 : -1);W dist3 = dist2 - (S[v] == '1' ? 1 : -1);tmp.push_back({dist2, X[c]});tmpc.push_back({dist2, X[c]});D[c].push_back(dist3);}if (res.size() < V.size()) res.swap(V);res.insert(res.end(), V.begin(), V.end());std::sort(tmpc.begin(), tmpc.end());FTc[v][i - C.G.begin(v)] = fenwick_tree_sparse<W, T>(tmpc);}std::sort(tmp.begin(), tmp.end());FT[v] = fenwick_tree_sparse<W, T>(tmp);return res;};dfs(dfs, root);}// vとの距離がl以上r未満の頂点の重みの和T sum(int v, W l, W r) {T res = 0;int p = v, c = -1;for (int i = 0; i < D[v].size(); i++, c = p, p = C.par[p]) {W dist = D[v][i];res += FT[p].prod(l - dist, r - dist);if (i) {int ord = O[c];res -= FTc[p][ord].prod(l - dist, r - dist);}}return res;}};int main() {io_init();int N;std::cin >> N;std::vector<std::pair<int, int>> E1;range(i, 0, N - 1) {int a, b;std::cin >> a >> b;a--, b--;E1.push_back({a, b});E1.push_back({b, a});}std::cin >> S;csr<int> G1(N, E1);std::vector<std::pair<int, weighted_edge<int>>> E2;auto f = [&](auto &&f, int v, int p) -> void {for (int i = G1.begin(v); i < G1.end(v); i++) {int to = G1.elist[i];if (to == p) continue;E2.push_back({v, {to, S[to] == '1' ? 1 : -1}});E2.push_back({to, {v, S[to] == '1' ? 1 : -1}});f(f, to, v);}};f(f, 0, -1);csr<weighted_edge<int>> G2(N, E2);hld<weighted_edge<int>> H(0, G2);centroid<weighted_edge<int>> C(G2);contour_sum_weighted<int, weighted_edge<int>> CS(H, C, std::vector<int>(N, 1));ll ans = 0;int cnt = 0;for (int i = 0; i < N; i++) {cnt += S[i] == '1';//std::cout << i << '\n';//for (int j = -3; j <= 3; j++) std::cout << CS.sum(i, j, j + 1) << ' ';//std::cout << '\n';ans += CS.sum(i, 1, 2 * N);}// std::cout << ans << " " << cnt << '\n';std::cout << (ans - cnt) / 2 + cnt << '\n';}