제출 #1162366

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1162366		hengliao	JOI tour (JOI24_joitour)	C++20	36 / 100	3093 ms	44072 KiB

joitour

#include "joitour.h"
#include<bits/stdc++.h>
using namespace std;

#define F first
#define S second
#define pll pair<ll, ll>
#define vll vector<ll>
#define pb push_back

typedef long long ll;

const ll mxN=2e5+5;

vll adj[mxN];
ll dp[mxN][5];
ll ans;
ll a[mxN];
ll sum[3];
ll con[mxN];
ll p[mxN];


void dfs(ll cur, ll par){
  for(ll i=0;i<5;i++) dp[cur][i]=0;
  dp[cur][a[cur]]++;
  sum[a[cur]]++;
  p[cur]=par;
  for(auto &chd:adj[cur]){
    if(chd==par) continue;
    dfs(chd, cur);
    for(ll i=0;i<3;i++){
      dp[cur][i]+=dp[chd][i];
    }
  }
}

void dfs3(ll cur, ll par){
  if(a[cur]==1){
    dp[cur][3]=dp[cur][0];
    dp[cur][4]=dp[cur][2];
  }
  for(auto &chd:adj[cur]){
    if(chd==par) continue;
    dfs3(chd, cur);
    for(ll i=3;i<5;i++){
      dp[cur][i]+=dp[chd][i];
    }
  }
}

void dfs2(ll cur, ll par){
  // if(a[cur]==1){
  //   for(auto &chd:adj[cur]){
  //     if(chd==par) continue;
  //     ans+=dp[chd][0]*(sum[2]-dp[chd][2]);
  //   }
  //   ans+=(sum[0]-dp[cur][0])*dp[cur][2];
  // }

  con[cur]=0;

  for(auto &chd:adj[cur]){
    if(chd==par) continue;
    dfs2(chd, cur);
    con[cur]+=dp[chd][3]*(dp[cur][2]-dp[chd][2]);
    con[cur]+=dp[chd][4]*(dp[cur][0]-dp[chd][0]);
    if(a[cur]==1) con[cur]+=dp[chd][0]*(dp[cur][2]-dp[chd][2]);
  }

  ans+=con[cur];

} 

void init(int n, vector<int> tep, vector<int> u, vector<int> v, int q) {
  for(ll i=0;i<n;i++){
    a[i]=tep[i];
  }
  for(ll i=0;i<n-1;i++){
    adj[u[i]].pb(v[i]);
    adj[v[i]].pb(u[i]);
  }
  ans=0;
  dfs(0, -1);
  dfs3(0, -1);
  dfs2(0, -1);
}

void change(int X, int Y) {
  a[X]=Y;
  ll cur=X;
  while(cur!=-1){
    ans-=con[cur];
    for(ll i=0;i<5;i++) dp[cur][i]=0;
    dp[cur][a[cur]]++;
    sum[a[cur]]++;
    for(auto &chd:adj[cur]){
      if(chd==p[cur]) continue;
      // dfs(chd, cur);
      for(ll i=0;i<3;i++){
        dp[cur][i]+=dp[chd][i];
      }
    }
    if(a[cur]==1){
      dp[cur][3]=dp[cur][0];
      dp[cur][4]=dp[cur][2];
    }
    for(auto &chd:adj[cur]){
      if(chd==p[cur]) continue;
      // dfs3(chd, cur);
      for(ll i=3;i<5;i++){
        dp[cur][i]+=dp[chd][i];
      }
    }
    con[cur]=0;

    for(auto &chd:adj[cur]){
      if(chd==p[cur]) continue;
      // dfs2(chd, cur);
      con[cur]+=dp[chd][3]*(dp[cur][2]-dp[chd][2]);
      con[cur]+=dp[chd][4]*(dp[cur][0]-dp[chd][0]);
      if(a[cur]==1) con[cur]+=dp[chd][0]*(dp[cur][2]-dp[chd][2]);
    }

    ans+=con[cur];
    cur=p[cur];
  }
}

long long num_tours() {
  return ans;
}

• Subtask 1: 6 / 6
• Subtask 2: 8 / 8
• Subtask 3: 6 / 6
• Subtask 4: 0 / 16
• Subtask 5: 16 / 16
• Subtask 6: 0 / 34
• Subtask 7: 0 / 14