AtCoder Begginer Contest #040 D - 道路の老朽化対策について

問題
D: 道路の老朽化対策について - AtCoder Beginner Contest 040 | AtCoder

N頂点、M辺の重み付き無向グラフが与えられる。
Q個のクエリがあり、初期位置v_iと辺を通れる境界のコストw_iが与えられる。
通れる境界のコストより大きいコストを持つ辺のみ、通ることが出来る。
各々のクエリについて、到達することの出来る頂点の数を答えよ。
1\le N\le 100000
0\le M\le 200000
1\le Q\le 100000


解法
あるノードから到達することの出来る頂点の数を数えるということから、あるノードから連結する部分グラフが境界のコストによって部分グラフの範囲が拡大していくようなイメージが持てる。グラフの連結性を次第に処理するのに効率的なデータ構造はUnionFindであるので、このデータ構造を使うことが思いつける。

w_iより大きい辺であれば到達可能、w_i以下の辺であれば到達不可能であるということは、w_iより大きい辺は全て連結したと見なすことが出来る。大きい辺から連結して頂点集合をまとめていくことは、UnionFindを用いたクラスカル法を用いることで処理できる。

クラスカル法は最小[最大]全域木を構築する手法ということに注目しがちだが、最小[最大]全域木を構築する過程で、貪欲に辺を処理して連結成分を拡大していくイメージが重要であると思う。

i番目のクエリのw_iより大きいコストをもつ辺のみを連結し、到達できる頂点集合の要素数を求めれば答えであるが、集合のサイズを求める実装はUnionFindに付属してライブラリ化しておくと良い。つまり、連結成分のある要素を与えた時、その連結成分のノードの数を求める実装をUnionFindに追加しておく(よく使う)。

ところで、クエリを入力順に処理しようとした時、例えば、w_0 = 1000で、w_1 = 2000であれば、0番目のクエリを処理しようとした段階で、1001以上のコストは全てつながっているような連結成分をUnionFindで構築していることが保証されている必要がある。そうすると、クエリ0を処理した後にクエリ1を処理することは出来ない。UnionFindはグラフの辺の連結はできても、切ることは出来ない。クエリを先読みしてコスト降順にソートしておけば、UnionFindの連結を切ることなしに、連結を繰り返しながら並行してクエリをコスト順に解決することが出来る。

class union_find {
  public:
    int par[MAX], rank[MAX], size[MAX], gnum;
    
    union_find(int N) {
        gnum = N;
        for (int i = 0; i < N; i++) {
            par[i] = i;
            rank[i] = 0;
            size[i] = 1;
        }
    }
 
    int find(int x)
    {
        if (par[x] == x) {
            return x;
        }
        return par[x] = find(par[x]);
    }
 
    void unite(int x, int y)
    {
        x = find(x);
        y = find(y);
        if (x == y) return;
        
        if (rank[x] < rank[y]) {
            par[x] = y;
            size[y] += size[x];
        } else {
            par[y] = x;
            size[x] += size[y];
            if (rank[x] == rank[y]) {
                rank[x]++;
            }
        }
        gnum--;
    }
 
    bool same(int x, int y)
    {
        return (find(x) == find(y));
    }
 
    int getSize(int x)
    {
        return size[find(x)];
    }
 
    int groups()
    {
        return gnum;
    }
};
 
typedef pair<int, int> pii;
 
int main() {
 
  int N, M; cin >> N >> M;
 
  typedef tuple<int, int, int> edge;
  vector<edge> edges;
  rep(i, M) {
    int a, b, y;
    cin >> a >> b >> y; a--, b--;
    edges.emplace_back(-y, a, b);
  }
 
  int Q; cin >> Q;
  vector<tuple<int, int, int>> query;
  rep(i, Q) {
    int v, w; cin >> v >> w; v--;
    query.emplace_back(-w, v, i);
  }
 
  sort(edges.begin(), edges.end());
  sort(query.begin(), query.end());
  rep(i, edges.size()) get<0>(edges[i]) *= -1;
  rep(i, query.size()) get<0>(query[i]) *= -1;
 
  union_find uf(N);
  vector<int> ans(Q, -inf);
 
  int k = 0;
 
  rep(i, edges.size()) {
    int eY, eA, eB; tie(eY, eA, eB) = edges[i];
 
    int qY = get<0>(query[k]);
    int qPos = get<1>(query[k]);
    int qIdx = get<2>(query[k]);
 
    if(qY < eY) {
      uf.unite(eA, eB);
    } else {
      ans[qIdx] = uf.getSize(qPos);
      i--;
      k++;
      if(k == query.size()) break;
    }
  }
  REP(i, k, query.size()) {
    int qPos = get<1>(query[i]);
    int qIdx = get<2>(query[i]);
    if(ans[qIdx] != -inf) assert(0);
    ans[qIdx] = uf.getSize(qPos);
  }
 
  rep(i, Q) {
    cout << ans[i] << endl;
  }
  
  return 0;
}