Submission #1000443

#TimeUsernameProblemLanguageResultExecution timeMemory
100044312345678JOI tour (JOI24_joitour)C++17
20 / 100
3097 ms38788 KiB
#include "joitour.h"
#include <bits/stdc++.h>

using namespace std;

#define ll long long

const int nx=2e5+5;

ll n, t[nx], cnt[3], dp[3][nx], res;
vector<int> d[nx];

void dfs(int u, int p)
{
    dp[0][u]=dp[1][u]=dp[2][u]=0;
    dp[t[u]][u]++;
    for (auto v:d[u]) if (v!=p) dfs(v, u), dp[0][u]+=dp[0][v], dp[1][u]+=dp[1][v], dp[2][u]+=dp[2][v];
    if (t[u]==1)
    {
        ll a=0, b=0;
        for (auto v:d[u])
        {
            if (v==p) continue;
            res+=dp[0][v]*b;
            res+=dp[2][v]*a;
            b+=dp[2][v];
            a+=dp[0][v];
        }
        //cout<<"here "<<u<<' '<<res<<'\n';
        res+=dp[0][u]*(cnt[2]-dp[2][u]);
        res+=dp[2][u]*(cnt[0]-dp[0][u]);
    }
}

void init(int N, std::vector<int> F, std::vector<int> U, std::vector<int> V, int Q) 
{
    n=N;
    for (int i=0; i<n; i++) t[i]=F[i], cnt[t[i]]++;
    for (int i=0; i<n-1; i++) d[U[i]].push_back(V[i]), d[V[i]].push_back(U[i]);
}

void change(int X, int Y) 
{
    cnt[t[X]]--;
    t[X]=Y;
    cnt[t[X]]++;
}

long long num_tours() 
{
    res=0;
    dfs(0, 0);
    return res;
}
#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...