BOJ 20171 - Dessert Café 링크
(2024.05.22 기준 G2)

문제

$n$개의 카페 후보지와 그들을 잇는 $n-1$개의 도로가 있다. 사이클은 존재하지 않는다.
$n$개의 후보지 중 $k$개의 후보지에는 아파트 단지가 있다.
각 후보지는 자신을 제외한 모든 후보지보다 가까운 아파트 단지가 하나 이상 있으면 그 후보지는 좋은 후보지가 된다. 이때 좋은 후보지의 개수 출력

알고리즘

그래프 탐색 문제다. 간선의 비용은 함정

풀이

일단 각 아파트 단지가 있는 후보지는 무조건 좋은 후보지다.

두 아파트 단지 $A$, $B$를 잇는 경로에 있는 후보지를 생각해보자. 그 경로에 있는 후보지들은 $A$와 따져보면 $B$와 가깝게 된다. 반대로 $B$와 따져보면 $A$와 가깝게 된다. 즉 간선의 비용에 상관없이 아파트 단지들을 잇는 경로의 모든 후보지가 전부 좋은 후보지가 된다.

코드

• C++

#include <bits/stdc++.h>
using namespace std;

typedef pair<int, int> pii;

const int MAXN = 100'001;

vector<pii> G[MAXN];
bool C[MAXN];

void dfs(int i, int p){
    for (auto [j, w]: G[i]){
        if (j == p) continue;
        dfs(j, i);
        if (C[j]) C[i] = true; // 아파트 단지를 거쳤다면 현재 i번 후보지도 좋은 후보지가 된다.
    }
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    int n, k; cin >> n >> k;

    for (int i, j, w, _ = 1; _ < n; _++){
        cin >> i >> j >> w;
        G[i].push_back({j, w});
        G[j].push_back({i, w});
    }

    memset(C, false, sizeof(C));
    for (int i; k; k--){
        cin >> i;
        C[i] = true;
    }

    // 아파트 단지가 있는 후보지에서부터 DFS 시작
    for (int i = 1; i <= n; i++) if (C[i]){
        dfs(i, 0);
        break;
    }

    int ans = 0;
    for (int i = 1; i <= n; i++) if (C[i]) ++ans;
    cout << ans;
}

• Python

import sys; input = sys.stdin.readline
sys.setrecursionlimit(111111)

def dfs(i, p):
    for j, w in G[i]:
        if j == p:
            continue
        dfs(j, i)
        if C[j]: # 아파트 단지를 거쳤다면 현재 i번 후보지도 좋은 후보지가 된다.
            C[i] = True

n, k = map(int, input().split())

G = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    i, j, w = map(int, input().split())
    G[i].append((j, w))
    G[j].append((i, w))

C = [False] * (n + 1)
for i in map(int, input().split()):
    C[i] = True

# 아파트 단지가 있는 후보지에서부터 DFS 시작
for i in range(1, n + 1):
    if C[i]:
        dfs(i, 0)
        break

print(C.count(True))