[백준 - 11725] 트리의 부모 찾기 - C++, Python

분류: 트리 union, 재귀, bfs, dfs /

 

문제

문제 링크

문제 설명

루트 없는 트리가 주어진다. 이때, 트리의 루트를 1이라고 정했을 때, 각 노드의 부모를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 노드의 개수 N (2 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에 트리 상에서 연결된 두 정점이 주어진다.

출력

첫째 줄부터 N-1개의 줄에 각 노드의 부모 노드 번호를 2번 노드부터 순서대로 출력한다.

 

풀이

- 트리 union (Python)

import sys
sys.setrecursionlimit(100_000)
read = sys.stdin.readline

N = int(read())
parent = [0] * (N + 1)
parent[1] = 1

def union(a, b):
    if parent[a] == 0:
        parent[a] = b
        if parent[b] == 0:
            parent[b] = a
    elif parent[b] == 0:
        parent[b] = a
    elif parent[b] != a:
        union(b, parent[b])
        parent[b] = a

ancestor = []
for _ in range(N-1):
    a, b = map(int, read().split())
    if a == 1:
        ancestor.append(b)
    elif b == 1:
        ancestor.append(a)
    else:
        union(a, b)

for anc in ancestor:
    union(1, anc)

for i in range(2, N + 1):
    print(parent[i])

`parent`: i노드의 부모 노드를 parent[i]에 저장. 트리는 부모가 하나이므로 배열에 저장할 수 있다.

`ancestor`: 조상 노드로 루트인 1과 연결된 노드들(루트의 직계 자손들).

 

알고리즘의 흐름은 다음과 같다.

1. 링크에 루트가 포함되어 있다면 루트와 연결된 노드를 조상 노드에 저장(ancestor)
2. 링크에 루트가 포함되지 않는다면 노드를 컴포넌트에 연결
   1. 링크에 연결된 두 노드가 모두 단일 노드이면 서로를 부모로 하는 사이클로 만들어 새 컴포넌트를 생성.
   2. 링크에 연결된 두 노드가 하나는 컴포넌트에 속해있고 다른 하나는 단일 노드라면 단일 노드를 컴포넌트에 연결.
   3. 각각 컴포넌트에 속해있다면 하나의 컴포넌트 전체를 다른 컴포넌트에 연결
3. 루트와 조상 노드 연결(루트는 자체로 컴포넌트(자기 참조)이므로 이 과정은 2-2혹은 2-3의 과정이다.)

2-2와 2-3에 의해 컴포넌트에는 항상 하나의 서로 참조하는 사이클을 갖으며 이 사이클은 컴포넌트의 루트다.

• 2-2에 의해 단일 노드는 사이클에 달린다.
• 2-3에 의해 컴포넌트끼리 합칠 때는 두 컴포넌트의 접점에서 종속되는 컴포넌트의 루트까지의 경로를 뒤집는다. 이때 종속되는 컴포넌트의 사이클(루트)은 사이클이 해제된다. 

# 예시

9 개의 노드에 대한 8개의 링크가 2-3, 1-2, 4-5, 6-7, 8-9, 8-7, 5-6, 4-3 순서로 입력되는 상황을 예로 들어보자.

먼저 첫 for문을 돌며 링크를 입력으로 받아 노드들을 연결해 컴포넌트를 만드는 과정이다.

+ 배열이 연속으로 붙어있는 것은 위에서부터 아래로 진행되는 과정으로 제일 위가 before, 제일 아래가 after다.

링크 1-2는 루트(1)를 포함하고 있으므로 ancestor에 2를 저장한다.

링크에 연결된 두 노드(a, b)가 모두 단일 노드라면(parent == 0) 서로 연결해 컴포넌트로 만들어준다.

 

각각 컴포넌트에 속해있다면 하나의 컴포넌트를 다른 컴포넌트에 연결한다.

union함수의 마지막 elif 부분이다. 재귀를 돌며 서로 참조하는 사이클을 찾고 찾은 사이클까지의 경로를 뒤집는다.

 

마지막으로 루트(1)와 조상 노드(ancestor)를 연결한다.

union 함수가 재귀적으로 호출되며 컴포넌트의 루트(사이클)를 찾고 (union - in)

컴포넌트의 루트를 찾고 나면 재귀를 멈추고 콜스택을 따라 나오며 사이클을 찾은 경로를 뒤집는다.(union - out)

 

- C++

#include <iostream>
#include <vector>
using namespace std;

int parent[100001];

void connect(int a, int b) {
    if(parent[a] && parent[b]) {
        if(parent[b] != a) {
            connect(b, parent[b]);
            parent[b] = a;
        }
        return;
    }
    if(!parent[b]) parent[b] = a;
    if(!parent[a]) parent[a] = b;
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int N; cin >> N;
    parent[1] = 1;
    vector<int> ancestor;
    for(int i=0; i<N-1; i++) {
        int a, b; cin >> a >> b;
        if(a == 1) ancestor.push_back(b);
        else if(b == 1) ancestor.push_back(a);
        else connect(a, b);
    }

    for(int x:ancestor) connect(1, x);

    for(int i=2; i<=N; i++)
        cout << parent[i] << '\n';
}

 

- bfs / dfs

from collections import deque
import sys
read = sys.stdin.readline

N = int(read())
tree = [0] * (N + 1)
tree[1] = 1
graph = [set() for _ in range(N + 1)]
for _ in range(N - 1):
    a, b = map(int, read().split())
    graph[a].add(b)
    graph[b].add(a)

q = deque([1])
while q:
    root = q.popleft()
    children = graph[root]
    for child in children:
        if tree[child] == 0:
            tree[child] = root
            q.append(child)

for i in range(2, N + 1):
    print(tree[i])

인접 리스트를 만들고 bfs를 돌며 루트(1)부터 퍼저나가며 자식을 매단다. dfs로도 구현할 수 있다.