結果

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

ソースコード

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

#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]-1
int 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) -1
int 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);
}
}
// pc
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 < rl > 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, blca(a, b)
// 1
// s-tlca(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;
}
}
// kx
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);
}
// vlr
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';
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0