Cod sursa(job #2924298)

Utilizator AlexandruBenescuAlexandru Benescu AlexandruBenescu Data 28 septembrie 2022 22:10:45
Problema Arbore Scor 0
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 2.44 kb
#include <bits/stdc++.h>
#define L 50005
#define COLORS 2
using namespace std;
ifstream fin("data.in");
ofstream fout("data.out");

vector <int> G[L];
bool col[L], vis[L];
int dp[L][COLORS], dp1[L][COLORS], dp2[L][COLORS], count_sub[L][COLORS], n, root = 1;

void DFS1(int node){
  int i;
  vis[node] = true;
  count_sub[node][col[node]] = 1;
  for (auto it : G[node])
    if (!vis[it]){
      DFS1(it);
      for (i = 0; i < COLORS; i++)
        count_sub[node][i] += count_sub[it][i];
    }
}

void DFS2(int node){
  int i;
  vis[node] = true;
  for (auto it : G[node])
    if (!vis[it]){
      DFS2(it);
      for (i = 0; i < COLORS; i++)
        dp1[node][i] += dp1[it][i] + count_sub[node][i] - (col[node] == i);
    }
}

void DFS3(int node){
  int i;
  vis[node] = true;
  for (i = 0; i < COLORS; i++)
    count_sub[node][i] = count_sub[root][i] - count_sub[node][i];
  for (auto it : G[node])
    if (!vis[it])
      DFS3(it);
}

void DFS4(int node){
  int i;
  vis[node] = true;
  for (auto it : G[node])
    if (!vis[it]){
      for (i = 0; i < COLORS; i++)
        dp2[it][i] = dp2[node][i] + count_sub[node][i];
      DFS4(it);
    }
}

void init_vis(){
  int i;
  for (i = 1; i <= n; i++)
    vis[i] = false;
}

void calc_dp(){
  int i, j;
  for (i = 1; i <= n; i++)
    for (j = 0; j < COLORS; j++)
      dp[i][j] = dp1[i][j] + dp2[i][j];
}

int main(){
  int q, t, node, i, j, a, b;
  fin >> n >> q;
  for (i = 1; i <= n; i++)
    fin >> col[i];
  for (i = 1; i < n; i++){
    fin >> a >> b;
    G[a].push_back(b);
    G[b].push_back(a);
  }
  DFS1(root);
  init_vis();
  DFS2(root);
  init_vis();
  DFS3(root);
  init_vis();
  DFS4(root);
  calc_dp();
  /**
  cout << "DP1:\n";
  for (i = 1; i <= n; i++){
    for (j = 0; j < COLORS; j++)
      cout << dp1[i][j] << " ";
    cout << "\n";
  }
  cout << "\n";
  cout << "DP2:\n";
  for (i = 1; i <= n; i++){
    for (j = 0; j < COLORS; j++)
      cout << dp2[i][j] << " ";
    cout << "\n";
  }
  cout << "\n";
  cout << "DP:\n";
  for (i = 1; i <= n; i++){
    for (j = 0; j < COLORS; j++)
      cout << dp[i][j] << " ";
    cout << "\n";
  }
  cout << "\n";
  **/
  for (j = 0; j < q; j++){
    fin >> t >> node;
    if (t == 1){
      if (col[i]){
        col[i] = false;
      }
      else{
        col[i] = true;
      }
    }
    else{
      fout << dp[node][col[node]] << "\n";
    }
  }
  return 0;
}