#include "joitour.h"
#include<bits/stdc++.h>
using namespace std;
#define ll long long int
#define MOD (1000000000+7)
#define MOD1 (998244353)
#define pb push_back
#define all(x) x.begin(), x.end()
#define en cout << '\n'
#define ff first
#define ss second
#define pii pair<int,int>
#define vi vector<int>
const int N = 2e5+100;
int tp[N], s[N];
array<ll, 5> dp[N]; // 0, 1, 2, 01, 21
ll ans;
vi g[N];
bitset<N> vis;
ll calc(int v){
return dp[v][2]*dp[v][3] + dp[v][0]*dp[v][4]; // 2*10 + 0*12
}
void pre(int v, int p){
s[v] = 1;
for(int u: g[v]){
if(!vis[u] && u != p){
pre(u, v);
s[v] += s[u];
}
}
}
int num;
int centro(int v, int p){
for(int u: g[v]){
if(!vis[u] && u != p && s[u] >= (num+1)/2) return centro(u, v);
}
return v;
}
void dfs(int v, int p, int c){
dp[v] = {0ll};
for(int u: g[v]){
if(u != p && !vis[u]){
dfs(u, v, c);
for(int j = 0; j < 6; ++j){
dp[v][j] += dp[u][j];
}
if(tp[v] == 1 && v != c) dp[v][3] += dp[u][0], dp[v][4] += dp[u][2];
}
}
if(v != c)
dp[v][tp[v]]++;
}
void f(int v){
pre(v, v);
num = s[v];
v = centro(v, v);
dfs(v, v, v);
ans += calc(v);
if(tp[v] == 1){
ans += dp[v][0] * dp[v][2];
}
for(int u: g[v]){
if(!vis[u]){
ans -= calc(u);
if(tp[v] == 1){
ans -= dp[u][0] * dp[u][2];
}
}
}
vis[v] = 1;
for(int u: g[v]){
if(!vis[u]) f(u);
}
}
void init(int n, std::vector<int> F, std::vector<int> U, std::vector<int> V, int Q) {
for(int i = 0; i + 1 < n; ++i){
g[U[i]].pb(V[i]);
g[V[i]].pb(U[i]);
}
for(int i = 1; i <= n; ++i){
tp[i - 1] = F[i - 1];
}
f(1);
}
void change(int X, int Y) {
assert(false);
}
long long num_tours() {
return ans;
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |