Codeforces Round #121 Div1 C. Fools and Roads

問題

n個の町が道路で結ばれている。グラフ(町, 道路)は木構造になっている。

以下の情報がk個与えられる。

u v

これは、ノードuとノードv間の単純路を人が移動したことを意味する。

各道路について、その道路を通った人の数を求めよ。

解法

u vという情報に対して、以下の処理を実行すればよい。

• uからrootまでのすべての枝に1を足す
• vからrootまでのすべての枝に1を足す
• wからrootまでのすべての枝から2を引く

ただし、wはLCA(u, v)とする。

毎回計算していたら遅いので、情報をすべて集めたあとに、葉から根方向にDPしながら和を計算すればよい。

実装

LCAのライブラリを作っておけば、張ってちょっと足して終わり。

#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#include <unordered_set>
#include <unordered_map>

using namespace std;

#define REP(i,n) for(int i=0; i<(int)(n); i++)
#define FOR(i,b,e) for (int i=(int)(b); i<(int)(e); i++)
#define ALL(x) (x).begin(), (x).end()

int n;
vector<int> edges[100000];
int par[100000];
int depth[100000];
long long cost[100000];
long long r[100000];
map<pair<int, int>, int> etoi;

class LCA {
  int V, logV;
  vector<int> depth;
  vector<vector<int> > parent;
  
  void build() {
    for (int k = 0; k + 1 < logV; k++) {
      for (int v = 0; v < V; v++) {
        if (parent[k][v] < 0) parent[k+1][v] = -1;
        else parent[k+1][v] = parent[k][parent[k][v]];
      }
    }
  }
public:
  // V: maximum number of nodes
  LCA(int V) {
    this->V = V;
    logV = 0;
    while (V > (1LL<<logV)) logV++;
    this->depth = vector<int>(V);
    this->parent = vector<vector<int> >(logV, vector<int>(V));
  }
  
  // N: number of nodes
  // p: parent of nodes (p[root] = -1)
  // d: depth of nodes (d[root = 0])
  void init(int N, int p[], int d[]) {
    for (int i = 0; i < N; i++) {
      parent[0][i] = p[i];
      depth[i] = d[i];
    }
    this->build();
  }
  
  // p: parent of nodes (p[root] = -1)
  // d: depth of nodes (d[root = 0])
  void init(vector<int> p, vector<int> d) {
    int N = d.size();
    for (int i = 0; i < N; i++) {
      parent[0][i] = p[i];
      depth[i] = d[i];
    }
    this->build();
  }
  
  int query(int u, int v) {
    if (depth[u] > depth[v]) swap(u, v);
    for (int k = 0; k < logV; k++) {
      if ((depth[v] - depth[u]) >> k & 1)
        v = parent[k][v];
    }
    if (u == v) return u;
    
    for (int k = logV-1; k >= 0; k--) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }
    return parent[0][u];
  }
};

void rec(int v, int p = -1) {
  for (auto w: edges[v]) {
    if (w != p) {
      rec(w, v);
      r[etoi[make_pair(v, w)]] += cost[w];
      cost[v] += cost[w];
    }
  }
}

void dfs(int v, int p = -1, int d = 0) {
  par[v] = p;
  depth[v] = d;
  for (auto &w: edges[v]) {
    if (w != p)
      dfs(w, v, d+1);
  }
}

int main() {
  scanf(" %d", &n);
  REP (i, n-1) {
    int x, y;
    scanf(" %d %d", &x, &y);
    --x, --y;
    etoi[make_pair(x, y)] = i;
    etoi[make_pair(y, x)] = i;
    edges[x].push_back(y);
    edges[y].push_back(x);
  }

  dfs(0);

  LCA lca(100000);
  lca.init(n, par, depth);

  int k;
  scanf(" %d", &k);
  REP (i, k) {
    int x, y;
    scanf(" %d %d", &x, &y);
    --x, --y;
    ++cost[x];
    ++cost[y];
    int z = lca.query(x, y);
    cost[z] -= 2;
  }

  rec(0);
  REP (i, n-1)
    cout << r[i] << " ";
  cout << endl;

  return 0;
}