test/aoj/other/2677.test.cpp

• View this file on GitHub
• Last update: 2022-08-17 04:19:08+09:00
• Problem: https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2677

Depends on

• 重軽分解 (data_structure/tree/heavy_light_decomposition.hpp)

Code

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2677"
#include <bits/stdc++.h>
using namespace std;
#include "../../../data_structure/tree/heavy_light_decomposition.hpp"
int main(){
  int n;
  cin >> n;
  vector<int> p(n);
  for (int i = 1; i < n; i++){
    cin >> p[i];
    p[i]--;
  }
  vector<vector<int>> c(n);
  for (int i = 1; i < n; i++){
    c[p[i]].push_back(i);
  }
  vector<int> bfs;
  queue<int> Q;
  Q.push(0);
  while (!Q.empty()){
    int v = Q.front();
    Q.pop();
    bfs.push_back(v);
    for (int w : c[v]){
      Q.push(w);
    }
  }
  heavy_light_decomposition T(p, c);
  long long ans = 0;
  for (int i = 0; i < n - 1; i++){
    ans += T.dist(bfs[i], bfs[i + 1]);
  }
  cout << ans << endl;
}

#line 1 "test/aoj/other/2677.test.cpp"
#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2677"
#include <bits/stdc++.h>
using namespace std;
#line 2 "data_structure/tree/heavy_light_decomposition.hpp"
/**
 * @brief 重軽分解
*/
struct heavy_light_decomposition{
  vector<int> p, sz, in, next;
  void dfs1(vector<vector<int>> &c, int v){
    sz[v] = 1;
    for (int &w : c[v]){
      dfs1(c, w);
      sz[v] += sz[w];
      if (sz[w] > sz[c[v][0]]){
        swap(w, c[v][0]);
      }
    }
  }
  void dfs2(vector<vector<int>> &c, int &t, int v){
    in[v] = t;
    t++;
    for (int w : c[v]){
      if (w == c[v][0]){
        next[w] = next[v];
      } else {
        next[w] = w;
      }
      dfs2(c, t, w);
    }
  }
  heavy_light_decomposition(vector<int> &p, vector<vector<int>> &c, int r = 0): p(p){
    int N = p.size();
    sz = vector<int>(N);
    dfs1(c, r);
    in = vector<int>(N);
    next = vector<int>(N, r);
    int t = 0;
    dfs2(c, t, r);
  }
  int lca(int u, int v){
    while (true){
      if (in[u] > in[v]){
        swap(u, v);
      }
      if (next[u] == next[v]){
        return u;
      }
      v = p[next[v]];
    }
  }
  int dist(int u, int v){
    int ans = 0;
    while (true){
      if (in[u] > in[v]){
        swap(u, v);
      }
      if (next[u] == next[v]){
        ans += in[v] - in[u];
        return ans;
      }
      ans += in[v] - in[next[v]] + 1;
      v = p[next[v]];
    }
  }
};
#line 5 "test/aoj/other/2677.test.cpp"
int main(){
  int n;
  cin >> n;
  vector<int> p(n);
  for (int i = 1; i < n; i++){
    cin >> p[i];
    p[i]--;
  }
  vector<vector<int>> c(n);
  for (int i = 1; i < n; i++){
    c[p[i]].push_back(i);
  }
  vector<int> bfs;
  queue<int> Q;
  Q.push(0);
  while (!Q.empty()){
    int v = Q.front();
    Q.pop();
    bfs.push_back(v);
    for (int w : c[v]){
      Q.push(w);
    }
  }
  heavy_light_decomposition T(p, c);
  long long ans = 0;
  for (int i = 0; i < n - 1; i++){
    ans += T.dist(bfs[i], bfs[i + 1]);
  }
  cout << ans << endl;
}

Back to top page