[백준] 11725번 : 트리의 부모 찾기 (DFS/BFS 실버Ⅱ) - Python, Java

문제

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

11725번: 트리의 부모 찾기

문제

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

입력

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

출력

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

예제 입력 1

7
1 6
6 3
3 5
4 1
2 4
4 7

예제 출력 1

4
6
1
3
1
4

예제 입력 2

12
1 2
1 3
2 4
3 5
3 6
4 7
4 8
5 9
5 10
6 11
6 12

예제 출력 2

1
1
2
3
3
4
4
5
5
6
6

접근

우선 예제 1번으로 문제를 이해해보면

1번부터 N번까지의 노드가 있다. 

입력받은 인접한 두 노드로 트리를 완성하고, 2번부터 N번까지 노드의 부모 노드가 무엇인지 차례로 출력하는 문제

이 문제는 1km 밖에서 물구나무 서서 봐도 BFS DFS 문제이기 때문에 이를 사용했다.

인접한 두 노드가 주어지므로 그래프를 딕셔너리로 나타낼 수 있고,

# 첫번째 예시에서

graph = {
    1: {4, 6}, 
    2: {4}, 
    3: {5, 6}, 
    4: {1, 2, 7}, 
    5: {3}, 
    6: {1, 3}, 
    7: {4}
    }

그럼 DFS 또는 BFS를 이용해 탐색을 진행하면서 부모 노드를 저장할 수 있다

물론 입력값이

1 6
6 3
3 5

와 같은 형태로 주어지므로 이를 그래프로 바꿔주는 작업을 진행해야 하는데, 처음에는 늘 그렇듯 파이썬 딕셔너리를 사용해 구현했다. 탐색을 진행할 때 그게 더 빠르기 때문에..

graph = {} # 그래프 딕셔너리 생성

for i in range(1, N+1):
    graph[i] = set()

for i in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].add(b)
    graph[b].add(a)

입력값의 형태 때문에 딕셔너리로 구현하기 위해서는 이러한 for 문을 두 개를 사용해야 하고 set()를 사용해야 했다.
이 부분에서 오히려 시간도, 메모리도 더 많이 사용되는 현상이 발생했다.
따라서 그래프를 다시 인접 리스트로 구현하는 것으로 변경해

graph = [[] for i in range(N+1)]

이 그래프를 사용하기로 했다.

이후엔 간단히 탐색을 구현하면 되는데, 유의할 점은 방문한 노드가 아니라 부모 노드를 반환해야 한다는 것.

따라서 단순히 방문한 노드를 저장하는 리스트(visited)를 생성해 반환하는 것이 아니라

1. 노드의 갯수를 길이로 갖는 리스트를 생성하고 (초기 요소 0)
2. 해당 노드를 인덱스로 삼고, 해당 인덱스 값에 부모 노드 값을 저장해준다.

풀이

Python 

먼저 DFS 구현이다.

DFS를 구현하는 방법은 두 가지, 재귀함수와 스택을 사용하는 방법이 있는데,

두 방법 모두 사용해서 풀어보았다.

DFS - 재귀함수

import sys
sys.setrecursionlimit(10**6)

N = int(sys.stdin.readline())

# 그래프 생성
graph = [[] for i in range(N+1)]
for i in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

# 방문 여부, 각 인덱스를 노드로 사용해 방문했으면 0을 지우고, 부모 노드를 저장한다.
visited = [0]*(N+1) 

'''dfs - 재귀함수로 구현'''
def dfs(node):
    for i in graph[node]: # 해당 노드에 인접한 노드 중에서
        if visited[i] == 0: # 방문한 적이 없다면
            visited[i] = node # 해당 노드를 부모 노드로 저장
            dfs(i)

dfs(1)
# 첫번째 예시라면 visited = [0, 4, 4, 6, 1, 3, 1, 4]

for x in range(2, N+1):
    print(visited[x]) # 각 인덱스(노드)에 저장된 부모 노드 가져오기

DFS - Stack

import sys

N = int(sys.stdin.readline())

# 그래프 생성
graph = [[] for i in range(N+1)]
for i in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

# 방문 여부, 각 인덱스를 노드로 사용해 방문했으면 0을 지우고, 부모 노드를 저장한다.
visited = [0]*(N+1) 

'''dfs - 스택으로 구현'''
def dfs(graph,node):

    stack = [node]
    while stack:
        top = stack.pop()
        for adj in graph[top]:
            if visited[adj] == 0: # 방문한 적이 없다면
                visited[adj] = top # 해당 노드를 부모 노드로 저장
                stack.append(adj)
    
    return visited

dfs(graph,1)
# 첫번째 예시라면 visited = [0, 4, 4, 6, 1, 3, 1, 4]

for x in range(2, N+1):
    print(visited[x]) # 각 인덱스(노드)에 저장된 부모 노드 가져오기

BFS

import sys
from collections import deque

N = int(sys.stdin.readline())

# 그래프 생성
graph = [[] for i in range(N+1)]
for i in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

# 방문 여부, 각 인덱스를 노드로 사용해 방문했으면 0을 지우고, 부모 노드를 저장한다.
visited = [0]*(N+1) 

""" BFS """
def bfs(start):
    deq = deque([start]) # 방문할 노드
    
    while deq:
        node = deq.popleft()
        
        for adj in graph[node]:
            if visited[adj] == 0:
                visited[adj] = node
                deq.append(adj)

bfs(1)
answer = visited[2:]
for x in answer:
    print(x)

각 풀이 별 메모리, 시간 비교

DFS / BFS 여부와 graph를 생성하는 방법(딕셔너리/인접리스트) 그리고 DFS 구현 방법(재귀, 스택) 에 따라 시간과 메모리는 다음과 같이 사용된다.

Java

DFS

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Stack;
import java.util.StringTokenizer;

public class Main {
    static int N;
    static ArrayList<Integer>[] graph;
    static int[] visited;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        N = Integer.parseInt(br.readLine());
        graph = new ArrayList[N + 1];
        visited = new int[N + 1];

        for (int i = 1; i <= N; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int i = 1; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            graph[a].add(b);
            graph[b].add(a);
        }

        dfs(1);

        for (int x = 2; x <= N; x++) {
            System.out.println(visited[x]);
        }
    }

    static void dfs(int node) {
        Stack<Integer> stack = new Stack<>();
        stack.push(node);

        while (!stack.isEmpty()) {
            int top = stack.pop();

            for (int adj : graph[top]) {
                if (visited[adj] == 0) {
                    visited[adj] = top;
                    stack.push(adj);
                }
            }
        }
    }
}

BFS

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.StringTokenizer;
import java.util.Queue;

public class Main {
    static int N;
    static ArrayList<Integer>[] graph;
    static int[] visited;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        N = Integer.parseInt(br.readLine());
        graph = new ArrayList[N+1];
        visited = new int[N+1];

        for (int i = 1; i <= N; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int i = 1; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            graph[a].add(b);
            graph[b].add(a);
        }
        bfs(1);

        for (int x = 2; x <= N; x++) {
            System.out.println(visited[x]);
        }
    }

    static void bfs(int start){
        Queue<Integer> deq = new ArrayDeque<>();
        deq.offer(start);

        while (!deq.isEmpty()) {
            int node = deq.poll();

            for (int adj : graph[node]) {
                if (visited[adj] == 0) {
                    visited[adj] = node;
                    deq.offer(adj);
                }
            }
        }
    }
}