Submission #1007766

#TimeUsernameProblemLanguageResultExecution timeMemory
1007766NamkhingJOI tour (JOI24_joitour)C++17
Compilation error
0 ms0 KiB
//#include "joitour.h"
#include <bits/stdc++.h>
#define ll long long
using namespace std;

ll n, q;
vector<ll> f, dp, order;
vector<vector<ll>> adj, cou;

void dfs(int u, int p) {
    for (ll& v : adj[u]) {
        if (v == p) {
            swap(v, adj[u].back());
            adj[u].pop_back();
            break;
        }
    }
    for (int v : adj[u]) {
        dfs(v, u);
    }   
    order.push_back(u);
}

void init(int N, std::vector<int> F, std::vector<int> U, std::vector<int> V, int Q) {
    n = N, q = Q;
    f.resize(N);
    dp.resize(N);
    cou.resize(3, vector<ll>(N));
    adj.resize(N);

    for (int i = 0; i < N; i++) {
        f[i] = F[i];
    }

    for (int j = 0; j < N - 1; j++) {
        adj[U[j]].push_back(V[j]);
        adj[V[j]].push_back(U[j]);
    }

    dfs(0, -1);
}

void change(int X, int Y) {
    f[X] = Y;
}

long long num_tours() {
    for (int u = 0; u < n; u++) {
        dp[u] = cou[0][u] = cou[1][u] = cou[2][u] = 0;
    }
    
    for (int u : order) {
        cou[f[u]][u] = 1;
        
        for (int v : adj[u]) {
            for (int i = 0; i < 3; i++) {
                cou[i][u] += cou[i][v];
            }
        }
    }

    ll sum = 0;
    
    for (int u : order) {
        if (f[u] == 1) {
            ll sum = 0;
            for (int v : adj[u]) {
                sum += cou[0][v] * cou[2][v];
            }
            
            dp[u] += cou[0][u] * cou[2][u] - sum;
            dp[u] += cou[0][u] * (cou[2][0] - cou[2][u]);
            dp[u] += cou[2][u] * (cou[0][0] - cou[0][u]);
        }

        for (int v : adj[u]) {
            dp[u] += dp[v];
        }
    }
    
    return dp[0];
}

int main() {
    int N;
    std::cin >> N;

    std::vector<int> F(N), U(N - 1), V(N - 1);

    for (int i = 0; i < N; i++) {
        std::cin >> F[i];
    }

    for (int j = 0; j < N - 1; j++) {
        std::cin >> U[j] >> V[j];
    }

    int Q;
    std::cin >> Q;

    std::vector<int> X(Q), Y(Q);

    for (int k = 0; k < Q; k++) {
        std::cin >> X[k] >> Y[k];
    }

    init(N, F, U, V, Q);

    std::cout << num_tours() << "\n";

    for (int k = 0; k < Q; k++) {
        change(X[k], Y[k]);
        std::cout << num_tours() << "\n";
    }
}

/*
3
0 1 2 
0 1 
1 2
0

3
0 1 2
0 1
1 2
2
2 0 
0 2

7 
1 0 2 2 0 1 0 
0 1
0 2
1 3
1 4
2 5
2 6
7
0 0
1 1
2 0
3 0
4 2
5 2
6 2
*/

Compilation message (stderr)

joitour.cpp: In function 'long long int num_tours()':
joitour.cpp:62:8: warning: unused variable 'sum' [-Wunused-variable]
   62 |     ll sum = 0;
      |        ^~~
/usr/bin/ld: /tmp/ccPMYkU9.o: in function `main':
stub.cpp:(.text.startup+0x0): multiple definition of `main'; /tmp/ccZr3Eb9.o:joitour.cpp:(.text.startup+0x0): first defined here
collect2: error: ld returned 1 exit status