ICPC国内予選2020-B: 追跡調査

問題

• https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1641

解説

与えられた $a_i,b_i$ の組についてそれぞれ, $a_i$ または $b_i$ が $p$ と連結ならば $a_i$ と $b_i$ の間に辺を張る.

頂点 $p$ の連結成分のサイズが求める答えとなる.

実装例

C++

• submission: https://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=6135818

#include <iostream>
#include <vector>

struct UnionFind {
  std::vector<int> parent_or_size;
  int cnt;

  UnionFind(int n) : parent_or_size(n, -1), cnt(n) {}

  void unite(int x, int y) {
    x = find_root(x);
    y = find_root(y);
    if (x == y) {
      return;
    }
    if (-parent_or_size[x] < -parent_or_size[y]) {
      std::swap(x, y);
    }
    parent_or_size[x] += parent_or_size[y];
    parent_or_size[y] = x;
    cnt--;
  }
  bool is_same_root(int x, int y) { return find_root(x) == find_root(y); }
  int find_root(int x) {
    if (parent_or_size[x] < 0) {
      return x;
    }
    return parent_or_size[x] = find_root(parent_or_size[x]);
  }
  int size(int x) { return -parent_or_size[find_root(x)]; }
};

void solve(int m, int n, int p) {
  p--;
  UnionFind uf_tree(m);
  for (int i = 0; i < n; i++) {
    int a, b;
    std::cin >> a >> b;
    a--, b--;
    if (uf_tree.is_same_root(a, p) || uf_tree.is_same_root(b, p)) {
      uf_tree.unite(a, b);
    }
  }
  std::cout << uf_tree.size(p) << "\n";
}

int main() {
  int m, n, p;

  while (true) {
    std::cin >> m >> n >> p;
    if (m == 0 && n == 0 && p == 0) {
      break;
    }

    solve(m, n, p);
  }

  return 0;
}

Python

• submission: https://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=6135832

#!/usr/bin/env python3


class UnionFind:
    def __init__(self, n):
        self.parent_or_size = [-1 for _ in range(n)]
        self.cnt = n

    def unite(self, x, y):
        x, y = self.find_root(x), self.find_root(y)
        if x == y:
            return
        if -self.parent_or_size[x] < -self.parent_or_size[y]:
            x, y = y, x
        self.parent_or_size[x] += self.parent_or_size[y]
        self.parent_or_size[y] = x
        self.cnt -= 1

    def is_same_root(self, x, y):
        return self.find_root(x) == self.find_root(y)

    def find_root(self, x):
        if self.parent_or_size[x] < 0:
            return x
        self.parent_or_size[x] = self.find_root(self.parent_or_size[x])
        return self.parent_or_size[x]

    def size(self, x):
        return -self.parent_or_size[self.find_root(x)]


while True:
    m, n, p = map(int, input().split())
    if m == 0 and n == 0 and p == 0:
        break

    p -= 1
    uf_tree = UnionFind(m)
    for _ in range(n):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        if uf_tree.is_same_root(a, p) or uf_tree.is_same_root(b, p):
            uf_tree.unite(a, b)

    print(uf_tree.size(p))