BOJ 25402 트리와 쿼리

https://www.acmicpc.net/problem/25402

25402번: 트리와 쿼리

 

매 쿼리마다 $S$가 주어졌을 때 $S$ 안의 정점들 주 연결되어 있는 정점 쌍의 개수를 찾는 문제입니다.

 

문제를 정리하면, $S$가 주어질 때 $S$안의 정점들이 몇 개의 컴포넌트로 나누어지는지, 그리고 각 컴포넌트에는 몇 개의 정점이 있는지 찾는 문제라고도 할 수 있습니다. 컴포넌트 내의 정점들은 모두 서로 연결되어 있다고 할 수 있으므로 매 컴포넌트마다 컴포넌트 안의 정점에서 두 개를 뽑는 경우의 수를 전부 더해주면 되기 때문입니다.

 

그렇다면 컴포넌트들을 어떻게 구할 수 있을까요? 나이브하게 생각해보면 매 쿼리당 dfs를 돌려서 구할 수 있습니다. 하지만 이 방법은 모든 간선을 확인해야 하기 때문에 쿼리당 $\mathcal{O}(n)$의 시간이 필요합니다. 시간을 줄이기 위해서는 컴포넌트에 속한 간선들만 따로 구할 수 있어야 합니다. 컴포넌트에 속한 간선들로만 새로 그래프를 만들어 dfs를 돌리면  $\mathcal{O}(k)$의 시간에 컴포넌트를 구할 수 있기 때문입니다.

 

컴포넌트에 속한 간선을 구하는 것은 나이브하게 생각해보면  $\mathcal{O}(k^2)$의 시간이 걸립니다. $S$내의 모든 정점 쌍에 대해 두 정점을 잇는 간선이 존재한다면 해당 간선이 컴포넌트에 속한다고 할 수 있기 때문입니다. 이를 최적화하기 위해 트리의 성질을 이용할 수 있습니다.

 

임의의 정점을 루트로 잡아 트리를 구성한다고 할 때, 정점 $a$와 정점 $b$가 이어져 있다는 것은 정점 $a$와 정점 $parent(a)$ 또는 정점 $b$와 정점 $parent(b)$가 연결되어 있다고 뜻입니다. 트리의 모든 간선은 부모 자식 관계의 정점들만을 이시습니다. 달리 말해 $S$안의 정점 $x$에 대하여 컴포넌트 내의 간선으로 사용될 수 있는 후보는 $x$와 $parent(x)$를 잇는 간선 뿐입니다. 따라서 $S$안의 모든 정점에 대해 자식과 부모가 모두 $S$안에 있는 경우의 간선들만 추가한 새로운 그래프를 만들고, 해당 그래프에서 dfs를 돌리면 답을 구할 수 있습니다.

 

아래는 코드입니다.

 

#include <bits/stdc++.h>
#define FastIO ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define endl '\n'
#define ends ' '
#define int long long
#define pii pair<int,int>

using namespace std;

int T, n, q, visited[250010], parent[250010], inS[250010];
vector<int> S;
vector<vector<int>> graph;
vector<vector<int>> newGraph;

int calc(int x) {
  return (x * (x - 1)) / 2;
}

int dfs(int x) {
  visited[x]++;
  int ret = 1;
  for (auto next : newGraph[x]) {
    if (visited[next]) continue;
    ret += dfs(next);
  }
  return ret;
}

void makeTree(int x) {
  visited[x]++;
  for (auto next : graph[x]) {
    if (visited[next]) continue;
    parent[next] = x;
    makeTree(next);
  }
}

int32_t main(void) {
  FastIO;
  T = 1;
  //cin >> T;

  while (T--) {
    cin >> n;

    graph.resize(n + 1);

    for (int i = 1; i < n; i++) {
      int a, b;
      cin >> a >> b;
      graph[a].push_back(b);
      graph[b].push_back(a);
    }

    makeTree(1);
    memset(visited, 0, sizeof(visited));

    cin >> q;

    while (q--) {
      int k;

      cin >> k;

      S.clear();
      S.resize(k);
      newGraph.clear();
      newGraph.resize(k + 1);

      for (int i = 0; i < k; i++) {
        cin >> S[i];
        inS[S[i]] = i + 1;
      }

      for (auto x : S) {
        if (inS[parent[x]]) {
          newGraph[inS[x]].push_back(inS[parent[x]]);
          newGraph[inS[parent[x]]].push_back(inS[x]);
        }
      }

      int ans = 0;

      for (int i = 0; i < k; i++) {
        if (visited[i]) continue;
        ans += calc(dfs(i));
      }

      cout << ans << endl;

      for (int i = 1; i <= k; i++) visited[i] = 0;
      for (auto x : S) inS[x] = 0;
    }
  }

  return 0;
}