『競プロ典型90問』ソースコード共有

q035.cpp

#include <bits/stdc++.h>
using namespace std;
using pii = pair<int, int>;
// QueryId: クエリId -> (Edges: 必要な辺数, Depth: 最小共通祖先(LCA)) のmap
using QED = unordered_map<int, pii>;
#define rep(i, n) for (int i=0; i<n; i++)

vector<int> depth;
vector<vector<int>> to, ids;

// 頂点xにおける部分木について, id毎に(必要な辺数, LCAの深さ)を返す
QED dfs(int x, int p) {
  QED qedx;
  // 行きがけは辺がない場合の解で初期化
  for (int id: ids[x]) {
    qedx.emplace(id, pii(0, depth[x]));
  }
  for (int y: to[x]) { // yはxの子
    if (y == p) continue;
    // 行きがけに子の深さを計算する
    depth[y] = depth[x] + 1;
    // 以上が行きがけ
    // それぞれの子の計算結果を取得する
    QED qedy = move(dfs(y, x));
    // 以下は戻りがけ
    // 子の方がサイズが大きい場合はswapしてならし計算数を200,000におさめる
    if (qedx.size() < qedy.size()) {
      swap(qedx, qedy);
    }
    // 以下のforループ内の処理はswapしたかどうかに関わらず同じ結果になる
    for (auto [id, _]: qedy) {
      if (qedx.find(id) != qedx.end()) { // 頂点xとyに共通のidがあり、ここで交わる
        // 既に確定していた必要辺数をマージ
        qedx[id].first += qedy[id].first;
        // xとyのLCAとの距離を必要辺数に加算する
        qedx[id].first += qedx[id].second - depth[x];
        qedx[id].first += qedy[id].second - depth[x];
        // 合流地点は現時点でのLCAなのでそのように更新
        qedx[id].second = depth[x];
        // cerr << id+1 << " met at " << x+1 << " or " << y+1 << " made (" << qedx[id].first << ", " << qedx[id].second << ")" << endl;
      } else {
        // idの合流は起こらないので単純に追加
        qedx.emplace(id, qedy[id]);
        // cerr << id+1 << " attended at " << x+1 << " or " << y+1 << endl;
      }
    }
  }
  return qedx;
}

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  int n;
  cin >> n;
  to.resize(n);
  rep(i, n-1) {
    int a, b;
    cin >> a >> b; a--; b--;
    to[a].push_back(b); to[b].push_back(a);
  }
  int q;
  cin >> q;
  ids.resize(n);
  rep(id, q) {
    int k;
    cin >> k;
    rep(j, k) {
      int v;
      cin >> v; v--;
      ids[v].push_back(id);
    }
  }
  depth.resize(n);
  QED res = dfs(0, -1); // 頂点0におけるid毎の(必要辺数, LCAの深さ)
  rep(i, q) {
    cout << res[i].first << endl;
  }
}

q035.rs

// original

use proconio::{input, marker::Usize1, fastout};
use std::collections::HashMap;

#[inline]
fn abs(x: usize, y: usize) -> usize { x.max(y)-x.min(y) }
#[derive(Debug, Default)]
struct Graph {
    n: usize,
    to: Vec<Vec<usize>>,
    nw: Vec<HashMap<usize, usize>>,
    dist: Vec<usize>,
    vm: VisitMap,
}
#[derive(Debug, Default)]
struct VisitMap {
    n: usize,
    discovered: Vec<bool>,
    to_visit: Vec<(usize, usize)>,
}
impl VisitMap {
    fn new(n: usize) -> Self {
        Self { n, discovered: vec![false; n], to_visit: vec![], }
    }
}
impl Graph {
    fn new(n: usize, ab: Vec<(usize, usize)>) -> Self {
        let mut to = vec![vec![]; n];
        for (a, b) in ab {
            to[a].push(b);
            to[b].push(a);
        }
        let nw = vec![HashMap::new(); n];
        Self { n, to, nw, ..Self::default() }
    }
    fn dfs(&mut self) {
        self.vm = VisitMap::new(self.n);
        self.dist = vec![0; self.n];
        let mut queue = vec![];
        queue.push(0);
        while let Some(x) = queue.pop() {
            self.vm.discovered[x] = true;
            for &y in &self.to[x] {
                if self.vm.discovered[y] { continue; }
                self.vm.to_visit.push( (y, x) );
                self.dist[y] = self.dist[x] + 1;
                queue.push(y);
            }
        }
    }
    fn dfs_post_order(&mut self, q: usize) -> Vec<usize> {
        let mut ans = vec![0; q];
        while let Some( (y, x) ) = self.vm.to_visit.pop() {
            if self.nw[x].len() < self.nw[y].len() {
                self.nw.swap(x, y);
            }
            let mut to_insert = Vec::<(usize, usize)>::new();
            for (rv, &d) in &self.nw[y] {
                if self.nw[x].contains_key(rv) {
                    ans[*rv] += self.nw[y][rv] - self.dist[x];
                    ans[*rv] += self.nw[x][rv] - self.dist[x];
                    to_insert.push((*rv, self.dist[x]));
                } else {
                    to_insert.push((*rv, d));
                }
            }
            for (v, d) in to_insert {
                self.nw[x].insert(v, d);
            }
        }
        ans
    }
}

#[fastout]
fn main() {
    input! {
        n: usize,
        ab: [(Usize1, Usize1); n-1],
        q: usize,
        vij: [[Usize1]; q],
    };
    let mut g = Graph::new(n, ab);
    g.dfs();
    let mut k = vec![0; q];
    for (i, vi) in vij.iter().enumerate() {
        k[i] = vi.len();
        for &v in vi {
            g.nw[v].insert(i, g.dist[v]);
        }
    }
    let ans = g.dfs_post_order(q);
    for x in ans {
        println!("{}", x);
    }
}

cppのコメントを一部変更