제출 #1160461

#제출 시각아이디문제언어결과실행 시간메모리
1160461mychecksedadJOI tour (JOI24_joitour)C++20
6 / 100
856 ms38012 KiB
#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 < 5; ++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];
  }else if(tp[v] == 0){
    ans += dp[v][4];
  }else{
    ans += dp[v][3];
  }
  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(0);
}

void change(int X, int Y) {
  assert(false);
}

long long num_tours() {
  return ans;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...