Submission #1199236

#	Time	Username	Problem	Language	Result	Execution time	Memory
1199236		Ronin13	JOI tour (JOI24_joitour)	C++20	100 / 100	2215 ms	307644 KiB

joitour

#include "joitour.h"
#include <bits/stdc++.h>
using ll = long long ;
using ull = unsigned ll;
#define f first 
#define s second 
#define pb push_back
#define epb emplace_back
using namespace std;
using pii = pair<int,int>;
using pll = pair<ll,ll>;

struct Data{

  vector <ll> f[3];
  int n;
  vector <ll> cnt;
  vector <vector<ll>>cnts;
  ll s_bad;
  Data() {
    cnt.resize(5);
  }
  void add(int pos, int val, int id) {
    while(pos < f[id].size()) {
      f[id][pos] += val;
      pos += (pos & -pos);
    }
  }
  ll get(int pos, int id) {
    ll res = 0;
    while(pos) {
      res += f[id][pos];
      pos -= (pos & -pos);
    }
    return res;
  }
  ll get(int l, int r, int id) {
    return get(r, id) - get(l - 1, id);
  }
};


const int nmax = 200001;
vector <vector <int>>vec(nmax);
vector <vector <int>>sz(nmax);
vector <vector<int>>in(nmax);
vector <vector <int>>g(nmax);
vector <int>y(nmax);
vector <vector<int>>lead_in(nmax);
Data dt[nmax];
vector <bool> blocked(nmax);
vector <int> siz(nmax);
ll ans;
int timer = 1;

vector <vector<ll>>cnt(nmax, vector<ll>(5));

void dfs(int v, int e) {
//  cout << v << ' ';
  siz[v] = 1;
  for(int to : g[v]) {
    if(to == e || blocked[to]) continue;
    dfs(to, v);
    siz[v] += siz[to];
  }
}

int find_cent(int v, int e, int n) {
  for(int to : g[v]) {
    if(blocked[to] || to == e) continue;
    if(siz[to] * 2 >= n) return find_cent(to, v, n);
  }
  return v;
}

void clear_dfs(int v, int e) {
  siz[v] = 0;
  for(int to : g[v]) {
    if(to == e || blocked[to]) continue;
    clear_dfs(to, v);
  }
  siz[v] = 0;
}


void update_cnt(int v, int to, bool x = false) {
  for(int i = 0; i < 5; i++) cnt[v][i] += cnt[to][i];
  if(y[v] == 1 && !x) {
    cnt[v][3] += cnt[to][0];
    cnt[v][4] += cnt[to][2];
  }
}

void DFS(int v, int e, int num, vector<int>&visited) {
    in[v].pb(timer++);
    lead_in[v].pb(num);
    sz[v].pb(1);
    visited.pb(v);
    for(int j = 0; j < 5; j++) cnt[v][j] = 0;
    for(int to : g[v]) {
      if(to == e || blocked[to]) continue;
      DFS(to, v, num, visited);
      update_cnt(v, to);
      sz[v].back() += sz[to].back();
    }
    cnt[v][y[v]]++;
}

void clearDFS(int v, int e) {
  cnt[v].assign(5, 0);
  for(int to : g[v]) {
    if(to == e || blocked[to]) continue;
    clearDFS(to, v);
  }
}

void decompose(int v) {
  dfs(v, v);
  int n = siz[v];
  int x = find_cent(v, v, siz[v]);

  clear_dfs(v, v);
  timer = 1;
  sz[x].pb(n);
  in[x].pb(1);
  vec[x].pb(x);
  lead_in[x].pb(0);
  int num = 0;
  vector <int> visited;
  for(int to : g[x]) {
    if(blocked[to]) continue;
    DFS(to, x, num, visited);
    update_cnt(x, to, true);
    dt[x].cnts.pb(cnt[to]);
    dt[x].s_bad += cnt[to][2] * cnt[to][0];
    num++;
  }

  cnt[x][y[x]]++;
  for(int j = 0; j< 3; j++) {
    dt[x].f[j].resize(n + 1);
  }
  dt[x].cnt = cnt[x];
  for(auto to : visited) {
    vec[to].pb(x);
    if(y[to] != 1) {
      dt[x].add(in[to].back(), 1, y[to]);
    }
    else {
      dt[x].add(in[to].back(), 1, 1);
      dt[x].add(in[to].back() + sz[to].back(), -1, y[to]);
    }
  }
  clearDFS(x, x);
  timer = 1;
  blocked[x] = true;
  for(int to : g[x]) {
      if(!blocked[to]) decompose(to);
  }
}

/*
ans += (dt[i].cnt[0] - to[0]) * to[4] + 
      (dt[i].cnt[2] - to[2]) * to[3] + 
      to[0] * (dt[i].cnt[4] - to[4]) + 
      to[2] * (dt[i].cnt[3] - to[3]);
        if(y[i] == 1) {
          ans += to[2] * (dt[i].cnt[0] - to[0]) + to[0] * (dt[i].cnt[2] - to[2]);
        }
        if(y[i] == 0) {
          ans += to[4];
        }
        if(y[i] == 2) {
          ans += to[3];
        } 
*/

void updCent0(int v, int id, ll sign) {
  ans += 2 * sign * (dt[v].cnt[4]); 
  dt[v].cnt[0] += sign;
}

void update0(int v, int id, ll sign) {
  int x = vec[v][id];
  
  int l = lead_in[v][id];
  vector <ll>&to = dt[x].cnts[l];
  if(v == x) {
    updCent0(v, id, sign);
    return;
  }
  int intm = in[v][id];
  int sztm = sz[v][id];
  ll cnt = dt[x].get(intm, 1);
  ans += 2 * sign * (dt[x].cnt[2] - to[2]) * cnt;
  ans += 2 * sign * (dt[x].cnt[4] - to[4]);
  if(y[x] == 1) {
    ans += 2 * sign * (dt[x].cnt[2] - to[2]);
  }
  to[0] += sign; dt[x].cnt[0] += sign;
  to[3] += sign * cnt;  dt[x].cnt[3] += cnt * sign;
  dt[x].s_bad += sign * to[2];
  dt[x].add(intm, sign, 0);
}

void updCent2(int v, int id, ll sign) {
  ans += 2 * sign * dt[v].cnt[3];
  dt[v].cnt[2] += sign;
}

void update2(int v, int id, ll sign) {
   int x = vec[v][id];
   
  if(v == x) {
    updCent2(v, id, sign);
    return;
  }
  int l = lead_in[v][id];
  vector <ll>&to = dt[x].cnts[l];
  int intm = in[v][id];
  int sztm = sz[v][id];
  ll cnt = dt[x].get(intm, 1);
  
  ans += 2 * sign * (dt[x].cnt[0] - to[0]) * cnt;
  ans += 2 * sign * (dt[x].cnt[3] - to[3]);
  
  if(y[x] == 1) {
    ans += 2 * sign * (dt[x].cnt[0] - to[0]);
  }
  to[2] += sign; dt[x].cnt[2] += sign;
  to[4] += sign * cnt;  dt[x].cnt[4] += sign * cnt;

  dt[x].s_bad += sign * to[0];
  
  dt[x].add(intm, sign, 2);
}

void updCent1(int v, ll sign) {
  ans += 2 * sign * (dt[v].cnt[0] * dt[v].cnt[2] - dt[v].s_bad);
  dt[v].cnt[1] += sign;
}

void update1(int v, int id, ll sign) {
   int x = vec[v][id];
  if(v == x) {
    updCent1(v, sign);
    return;
  }
  int l = lead_in[v][id];
  vector <ll>&to = dt[x].cnts[l];
  int intm = in[v][id];
  int sztm = sz[v][id];
 //cout << x << ' ' << intm << ' ' << sztm << ' ';
  ll num0 = dt[x].get(intm, intm + sztm - 1, 0);
 // cout << num0 << ' ' ;
  ll num2 = dt[x].get(intm, intm + sztm - 1, 2);
 // cout << num2 << '\n';
  ans += sign * 2 * num0 * (dt[x].cnt[2] - to[2]);
  ans += sign * 2 * num2 * (dt[x].cnt[0] - to[0]);
  to[1] += sign; to[3] += num0 * sign; to[4] += num2 * sign;
  dt[x].cnt[1] += sign; dt[x].cnt[3] += num0 * sign; dt[x].cnt[4] += num2 * sign;  
 
     dt[x].add(intm, sign, 1);
  dt[x].add(intm + sztm, -sign, 1);
  
}


void init(int n, std::vector<int> f, std::vector<int> u, std::vector<int> v,int q) {
    for(int i = 0; i < f.size(); i++) {
      y[i] = f[i];
    }
   // cout << 1;
    for(int i = 0; i < u.size(); i++) {
      g[u[i]].pb(v[i]);
      g[v[i]].pb(u[i]);
    }
   // cout << 1;
    decompose(0);
   // cout << 1;
    for(int i = 0; i < n; i++) {
     
      for(auto to : dt[i].cnts) {
        ans += (dt[i].cnt[0] - to[0]) * to[4] + (dt[i].cnt[2] - to[2]) * to[3] + to[0] * (dt[i].cnt[4] - to[4]) + to[2] * (dt[i].cnt[3] - to[3]);
        if(y[i] == 1) {
          ans += to[2] * (dt[i].cnt[0] - to[0]) + to[0] * (dt[i].cnt[2] - to[2]);
        }
        if(y[i] == 0) {
          ans += to[4];
        }
        if(y[i] == 2) {
          ans += to[3];
        } 
      }
     // cout << ans << ' ';
      
    }
}

void change(int X, int Y) {
  for(int i = 0; i < vec[X].size(); i++) {
    if(y[X] == 0) {
      update0(X, i, -1);
    }
    if(y[X] == 1)
    update1(X, i, -1);

    if(y[X] == 2) 
   update2(X, i, -1);
  }
 // return;
  y[X] = Y;
  for(int i = 0; i < vec[X].size(); i++) {
    if(y[X] == 0) {
      update0(X, i, 1);
    }
    
    if(y[X] == 1)
    update1(X, i, 1);
    if(y[X] == 2) 
    update2(X, i, 1);
  }
}

long long num_tours() {
  
  return ans / 2;
}

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