結果
問題 | No.2892 Lime and Karin |
ユーザー |
![]() |
提出日時 | 2024-09-13 23:07:32 |
言語 | C++17(gcc12) (gcc 12.3.0 + boost 1.87.0) |
結果 |
WA
|
実行時間 | - |
コード長 | 18,541 bytes |
コンパイル時間 | 3,630 ms |
コンパイル使用メモリ | 247,012 KB |
実行使用メモリ | 55,316 KB |
最終ジャッジ日時 | 2024-09-13 23:07:45 |
合計ジャッジ時間 | 12,290 ms |
ジャッジサーバーID (参考情報) |
judge1 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 16 WA * 36 |
ソースコード
/*#pragma GCC target("avx2")#pragma GCC optimize("O3")#pragma GCC optimize("unroll-loops")//*/// #include <atcoder/all>#include <bits/stdc++.h>using namespace std;// using namespace atcoder;// #define _GLIBCXX_DEBUG#define DEBUG(x) cerr << #x << ": " << x << endl;#define DEBUG_VEC(v) \cerr << #v << ":"; \for (int iiiiiiii = 0; iiiiiiii < v.size(); iiiiiiii++) \cerr << " " << v[iiiiiiii]; \cerr << endl;#define DEBUG_MAT(v) \cerr << #v << endl; \for (int iv = 0; iv < v.size(); iv++) { \for (int jv = 0; jv < v[iv].size(); jv++) { \cerr << v[iv][jv] << " "; \} \cerr << endl; \}typedef long long ll;// #define int ll#define vi vector<int>#define vl vector<ll>#define vii vector<vector<int>>#define vll vector<vector<ll>>#define pii pair<int, int>#define pis pair<int, string>#define psi pair<string, int>#define pll pair<ll, ll>template <class S, class T>pair<S, T> operator+(const pair<S, T> &s, const pair<S, T> &t) {return pair<S, T>(s.first + t.first, s.second + t.second);}template <class S, class T>pair<S, T> operator-(const pair<S, T> &s, const pair<S, T> &t) { return pair<S, T>(s.first - t.first, s.second - t.second); }template <class S, class T>ostream &operator<<(ostream &os, pair<S, T> p) {os << "(" << p.first << ", " << p.second << ")";return os;}#define rep(i, n) for (int i = 0; i < (int)(n); i++)#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)#define rrep(i, n) for (int i = (int)(n) - 1; i >= 0; i--)#define rrep1(i, n) for (int i = (int)(n); i > 0; i--)#define REP(i, a, b) for (int i = a; i < b; i++)#define in(x, a, b) (a <= x && x < b)#define all(c) c.begin(), c.end()void YES(bool t = true) {cout << (t ? "YES" : "NO") << endl;}void Yes(bool t = true) { cout << (t ? "Yes" : "No") << endl; }void yes(bool t = true) { cout << (t ? "yes" : "no") << endl; }void NO(bool t = true) { cout << (t ? "NO" : "YES") << endl; }void No(bool t = true) { cout << (t ? "No" : "Yes") << endl; }void no(bool t = true) { cout << (t ? "no" : "yes") << endl; }template <class T>bool chmax(T &a, const T &b) {if (a < b) {a = b;return 1;}return 0;}template <class T>bool chmin(T &a, const T &b) {if (a > b) {a = b;return 1;}return 0;}template <class T, class U>T ceil_div(T a, U b) {return (a + b - 1) / b;}template <class T>T safe_mul(T a, T b) {if (a == 0 || b == 0) return 0;if (numeric_limits<T>::max() / a < b) return numeric_limits<T>::max();return a * b;}#define UNIQUE(v) v.erase(std::unique(v.begin(), v.end()), v.end());const ll inf = 1000000001;const ll INF = (ll)1e18 + 1;const long double pi = 3.1415926535897932384626433832795028841971L;int popcount(ll t) { return __builtin_popcountll(t); }vector<int> gen_perm(int n) {vector<int> ret(n);iota(all(ret), 0);return ret;}// int dx[4] = {1, 0, -1, 0}, dy[4] = {0, 1, 0, -1};// int dx2[8] = { 1,1,0,-1,-1,-1,0,1 }, dy2[8] = { 0,1,1,1,0,-1,-1,-1 };vi dx = {0, 0, -1, 1}, dy = {-1, 1, 0, 0};vi dx2 = {1, 1, 0, -1, -1, -1, 0, 1}, dy2 = {0, 1, 1, 1, 0, -1, -1, -1};struct Setup_io {Setup_io() {ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);cout << fixed << setprecision(25);cerr << fixed << setprecision(25);}} setup_io;// constexpr ll MOD = 1000000007;constexpr ll MOD = 998244353;// #define mp make_pairenum Type { Vertex,Compress,Rake,AddEdge,AddVertex };struct StaticTopTree {// https://atcoder.jp/contests/abc351/editorial/9868// https://atcoder.jp/contests/abc351/submissions/52777033 を元に、自分好みに書き方を変えたりコメントを追加したり辺クエリに対応したりした// 全方位木 DP のライブラリと考え方が似ているのでそちらも参照すると良い// G は n 頂点の木を表す連結リストであるとする。// 木を HLD のように heavy path と light path に分解する。この時、親からの辺が light edge である頂点と heavy edgeである頂点の二種類に分類できる。// 前者の頂点を軽頂点、後者の頂点を重頂点と呼ぶ。ただし、根は重頂点であるとする。// 頂点も辺も存在しない状態から、以下の操作を行うことで元の木を構築することを考える。この際、頂点が各操作を表すような操作列を表す木も同時に構築する。//// A-1. 子に軽頂点を持たない頂点を追加する(Vertex)// B. 軽頂点が根のクラスタ(部分木的なもの)に親辺を追加する。親頂点は追加しないことに注意(AddEdge)// C. 複数の「B. で作られたもの」をマージする(Rake)// A-2. 「C. で作られたもの」の親頂点を追加し、「C. で作られたもの」とくっつける(AddVertex)// D. 重頂点とその親頂点をマージすることで heavy path を伸ばす(Compress)//// 実装上は C, D は二つのクラスタのマージを繰り返すことで表現するため、「操作列を表す木」は二分木となる。// この際、マージする順番を工夫することで完全二分木に近づけることができ、木の深さを O(log n) に抑えることができる。vector<vector<pair<int, int>>> G; // 元のグラフ、(to, edge_idx) のリストで表され、親への辺は含まれないint root; // 元のグラフでの根vector<int> par_edge; // 元のグラフで各頂点の親への辺のインデックスint stt_root; // 操作列を表す木の根vector<int> P, L, R; // 操作列を表す木における、親、左の子、右の子vector<Type> T; // 操作列を表す木における、各頂点の操作タイプvector<int> merged_edge; // compress / add_edge の操作を表す頂点において、その辺のインデックスint add_buf; // 操作列を表す木における、新たに追加する頂点のインデックスStaticTopTree(const vector<vector<pair<int, int>>> &_G = {}, int _root = 0) : G(_G), root(_root) {int n = G.size();if (n == 0) {return;}par_edge.resize(n, -1);P.resize(4 * n, -1), L.resize(4 * n, -1), R.resize(4 * n, -1);T.resize(4 * n, Type::Vertex);merged_edge.resize(4 * n, -1);add_buf = n;build();}int size() const {// 操作列を表す木のサイズを返すreturn add_buf;}private:int dfs(int c, int par = -1) {// 部分木のサイズを計算し、最も大きい部分木をリストの先頭に持ってくる// また、各頂点の親への辺のインデックスを記録するint s = 1, best = 0;for (int i = 0; i < (int)G[c].size(); i++) {auto [d, ei] = G[c][i];par_edge[d] = ei;int t = dfs(d, c);s += t;if (best < t) best = t, swap(G[c][i], G[c][0]);}return s;}int add(int k, int l, int r, Type t) {if (k == -1) k = add_buf++;P[k] = -1, L[k] = l, R[k] = r, T[k] = t;if (l != -1) P[l] = k;if (r != -1) P[r] = k;return k;}// 以下の pair<int, int> を返す関数は、(操作に対応する頂点のインデックス, 部分木に含まれる実頂点の数) を返す// ただし、操作に対応する頂点が作られない場合は (-1, 0) を返すpair<int, int> merge(const vector<pair<int, int>> &a, Type t) {assert(a.size() > 0);if (a.size() == 1) return a[0];int u = 0;for (auto &[_, s] : a)u += s;// 二つに振り分けたあとにマージするvector<pair<int, int>> b, c;for (auto &[i, s] : a)(u > s ? b : c).emplace_back(i, s), u -= s * 2;auto [i, si] = merge(b, t);auto [j, sj] = merge(c, t);auto ret = pair<int, int>(add(-1, i, j, t), si + sj);if (t == Type::Compress) {// compress の操作を表す頂点において、その辺のインデックスを記録する// c[0] が今回マージされる子頂点であり、c[0] が Vertex, AddVertex であることは `compress` 関数の仕様から保証されるmerged_edge[ret.first] = par_edge[c[0].first];}return ret;}pair<int, int> compress(int i) {vector<pair<int, int>> chs;while (true) {chs.push_back(add_vertex(i));if (G[i].size() == 0) {break;}i = G[i][0].first;}return merge(chs, Type::Compress);}pair<int, int> rake(int i) {vector<pair<int, int>> chs;for (int j = 1; j < (int)G[i].size(); j++)chs.push_back(add_edge(G[i][j].first));return chs.empty() ? make_pair(-1, 0) : merge(chs, Type::Rake);}pair<int, int> add_edge(int i) {// i の親は light path であることが保証されるauto [j, sj] = compress(i);auto ret = pair<int, int>(add(-1, j, -1, Type::AddEdge), sj);// add_edge の操作を表す頂点において、その辺のインデックスを記録する// i が Vertex, AddVertex であることは `rake` 関数の仕様から保証されるmerged_edge[ret.first] = par_edge[i];return ret;}pair<int, int> add_vertex(int i) {auto [j, sj] = rake(i);return {add(i, j, -1, j == -1 ? Type::Vertex : Type::AddVertex), sj + 1};}void build() {dfs(root);auto [i, n] = compress(root);stt_root = i;}};template <class E, class V>struct StaticTopTreeDp {// static top tree を用いて以下を実現する// - 各頂点や辺がある値を持つ木に関する木 DP を行う// - 各頂点や辺が持つ値が変わった時に O(log n) で木 DP の値を更新する// - https://atcoder.jp/contests/abc351/submissions/52951833// - https://judge.yosupo.jp/submission/205112// - 部分木のサイズ程度の長さの列を値に持つ DP は、分割統治的な考え方で O(n (log n)^2) に高速化できる// - https://atcoder.jp/contests/abc269/submissions/52956082//// 全方位木 DP のライブラリと考え方が似ているのでそちらも参照すると良い// E: クラスタに親辺を追加したもの、それらをマージしたものが持つべき答えの型。AddEdge, Rake に対応。// V: 頂点を根に持つクラスタが持つべき答えの型。Vertex, AddVertex, Compress に対応。// ここで、AddVertex と Compress に対応するため、重頂点の子の情報は持たないことに注意。StaticTopTree stt;function<V(int)> vertex; // 頂点のインデックスfunction<E(V, int)> add_edge; // (クラスタの答え, 辺のインデックス)function<E(E, E)> rake; // (マージ対象の左クラスタの答え, 右)function<V(E, int)> add_vertex; // (クラスタの答え, 頂点のインデックス)function<V(V, V, int)> compress; // (マージ対象の親クラスタの答え, 子, 辺のインデックス)vector<int> child; // 辺ごとに、子の方の頂点を記録するvector<E> dp_e; // stt の各頂点に対応するクラスタの答え。頂点が AddEdge, Rake でない場合は使用されない想定。vector<V> dp_v; // stt の各頂点に対応するクラスタの答え。頂点が Vertex, AddVertex, Compress でない場合は使用されない想定。StaticTopTreeDp(vector<vector<pair<int, int>>> G,function<V(int)> vertex,function<E(V, int)> add_edge,function<E(E, E)> rake,function<V(E, int)> add_vertex,function<V(V, V, int)> compress,int root = 0): vertex(vertex), add_edge(add_edge), rake(rake), add_vertex(add_vertex), compress(compress) {child.resize(G.size() - 1, -1);preprocess_graph(G, root, -1);stt = StaticTopTree(G, root);dp_e.resize(stt.size());dp_v.resize(stt.size());}V fill_dp_all() {// DP を全て埋める。最初に一度だけ呼ばれる想定。dfs(stt.stt_root);return dp_v[stt.stt_root];}V update_vertex(int u) {// 頂点 u の値が変わった後に、その影響を木 DP に反映する。O(log n)while (u != -1) {update(u);u = stt.P[u];}return dp_v[stt.stt_root];}V update_edge(int e) {// 辺 e の値が変わった後に、その影響を木 DP に反映する。O(log n)return update_vertex(child[e]);}private:void preprocess_graph(vector<vector<pair<int, int>>> &G, int now, int par) {int par_idx = -1;for (int i = 0; i < (int)G[now].size(); i++) {if (G[now][i].first == par) {par_idx = i;} else {child[G[now][i].second] = G[now][i].first;preprocess_graph(G, G[now][i].first, now);}}if (par_idx != -1) {swap(G[now][par_idx], G[now].back());G[now].pop_back();}}void update(int u) {if (stt.T[u] == Type::Vertex) {dp_v[u] = vertex(u);} else if (stt.T[u] == Type::Compress) {dp_v[u] = compress(dp_v[stt.L[u]], dp_v[stt.R[u]], stt.merged_edge[u]);} else if (stt.T[u] == Type::Rake) {dp_e[u] = rake(dp_e[stt.L[u]], dp_e[stt.R[u]]);} else if (stt.T[u] == Type::AddEdge) {dp_e[u] = add_edge(dp_v[stt.L[u]], stt.merged_edge[u]);} else {dp_v[u] = add_vertex(dp_e[stt.L[u]], u);}}void dfs(int now) {if (stt.L[now] != -1) dfs(stt.L[now]);if (stt.R[now] != -1) dfs(stt.R[now]);update(now);}};vector<vector<pii>> G;vector<pii> es;vi a;struct V {int root;vector<ll> num;int offset;V() {}V(int root, vector<ll> num, int offset) : root(root), num(num), offset(offset) {}};struct E {int root;vector<ll> num;int offset;E() {}E(int root, vector<ll> num, int offset) : root(root), num(num), offset(offset) {}};vl add(vl a, vl b) {if (a.size() < b.size()) {swap(a, b);}for (int i = 0; i < b.size(); i++) {a[i] += b[i];}return a;}ll ans = 0;V vertex(int i) {if (a[i] == 1) {ans++;}return V(i, {1}, a[i]);}E add_edge(V v, int ei) {auto [a, b] = es[ei];assert(v.root == a || v.root == b);int r;if (v.root == a) {r = b;} else {r = a;}return E(r, v.num, v.offset);}E rake(E l, E r) {assert(l.root == r.root);// DEBUG(l.root);// DEBUG_VEC(l.num);// DEBUG(l.offset);// DEBUG(r.root);// DEBUG_VEC(r.num);// DEBUG(r.offset);ll add = 0;int def = a[l.root];rep(i, l.num.size()) {rep(j, r.num.size()) {if (i + l.offset + j + r.offset + def > 0) {add += l.num[i] * r.num[j];}}}// DEBUG(add);ans += add;if (l.offset > r.offset) {swap(l, r);}for (int i = 0; i < r.num.size(); i++) {int ori_idx = i + r.offset;int new_idx = ori_idx - l.offset;while (l.num.size() <= new_idx) {l.num.push_back(0);}l.num[new_idx] += r.num[i];}return E(l.root, l.num, l.offset);}V add_vertex(E d, int i) {assert(d.root == i);if (a[i] == 1) {ans++;}// DEBUG(d.root);// DEBUG_VEC(d.num);// DEBUG(d.offset);ll add = 0;rep(j, d.num.size()) {if (j + d.offset + a[i] > 0) {add += d.num[j];}}// DEBUG(add);ans += add;vl num1 = d.num, num2 = {1};int offset1 = d.offset + a[i], offset2 = a[i];if (offset1 > offset2) {swap(num1, num2);swap(offset1, offset2);}for (int i = 0; i < num2.size(); i++) {int ori_idx = i + offset2;int new_idx = ori_idx - offset1;while (num1.size() <= new_idx) {num1.push_back(0);}num1[new_idx] += num2[i];}return V(i, num1, offset1);};V compress(V p, V c, int ei) {// DEBUG(p.root);// DEBUG_VEC(p.num);// DEBUG(p.offset);// DEBUG(c.root);// DEBUG_VEC(c.num);// DEBUG(c.offset);ll add = 0;rep(i, p.num.size()) {rep(j, c.num.size()) {if (i + p.offset + j + c.offset > 0) {add += p.num[i] * c.num[j];}}}// DEBUG(add);ans += add;c.offset += a[p.root];int root = p.root;if (p.offset > c.offset) {swap(p, c);}for (int i = 0; i < c.num.size(); i++) {int ori_idx = i + c.offset;int new_idx = ori_idx - p.offset;while (p.num.size() <= new_idx) {p.num.push_back(0);}p.num[new_idx] += c.num[i];}return V(root, p.num, p.offset);}signed main() {int n;cin >> n;G.resize(n);rep(i, n - 1) {int a, b;cin >> a >> b;a--, b--;G[a].push_back({b, i});G[b].push_back({a, i});es.emplace_back(a, b);}string s;cin >> s;rep(i, n) {if (s[i] == '0') {a.push_back(-1);} else {a.push_back(1);}}StaticTopTreeDp<E, V> sttd(G, vertex, add_edge, rake, add_vertex, compress);sttd.fill_dp_all();cout << ans << endl;}