結果

問題 No.2892 Lime and Karin
ユーザー fuppy_kyoprofuppy_kyopro
提出日時 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
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

/*
#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_pair
enum 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;
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0