백준 1325: 효율적인 해킹

관련 링크 : https://www.acmicpc.net/problem/1325

백준 1325 : 효율적인 해킹

해커 김지민은 잘 알려진 어느 회사를 해킹하려고 한다. 이 회사는 N개의 컴퓨터로 이루어져 있다. 김지민은 귀찮기 때문에, 한 번의 해킹으로 여러 개의 컴퓨터를 해킹 할 수 있는 컴퓨터를 해킹하려고 한다.

이 회사의 컴퓨터는 신뢰하는 관계와, 신뢰하지 않는 관계로 이루어져 있는데, A가 B를 신뢰하는 경우에는 B를 해킹하면, A도 해킹할 수 있다는 소리다.

이 회사의 컴퓨터의 신뢰하는 관계가 주어졌을 때, 한 번에 가장 많은 컴퓨터를 해킹할 수 있는 컴퓨터의 번호를 출력하는 프로그램을 작성하시오.

입력

첫째 줄에, N과 M이 들어온다. N은 10,000보다 작거나 같은 자연수, M은 100,000보다 작거나 같은 자연수이다. 둘째 줄부터 M개의 줄에 신뢰하는 관계가 A B와 같은 형식으로 들어오며, “A가 B를 신뢰한다”를 의미한다. 컴퓨터는 1번부터 N번까지 번호가 하나씩 매겨져 있다.

5 4
3 1
3 2
4 3
5 3

출력

첫째 줄에, 김지민이 한 번에 가장 많은 컴퓨터를 해킹할 수 있는 컴퓨터의 번호를 오름차순으로 출력한다.

1 2

풀이

이때까지 모든 문제를 DFS로 먼저 접근해서 풀었고, 이 문제도 동일하게 접근했다. 앞선 방식들과 차이라면 이번 문제는 단 방향 그래프라는 점이다.

import sys

sys.setrecursionlimit(10000)
input = sys.stdin.readline

def getInput():
    N, M = map(int, input().strip().split())
    relations = [[] for _ in range(N+1)]
    for _ in range(M):
        end, start = map(int, input().strip().split())
        relations[start].append(end) 
    return N, relations
    
def dfs(start, relations, visited):
    for end in relations[start]:
        if not visited[end]: #방문하지 않은 경우만 진행
            visited[end] = 1
            dfs(end, relations, visited)

def solution1():
    N, relations = getInput()
    max_val = 0
    answers = []
    for node in range(1, N+1): #모든 노드에서 컴퓨터 개수 세기
        visited = [0 for _ in range(N+1)]
        visited[node] = 1
        dfs(node, relations, visited)
        val = sum(visited)
        if val > max_val:
            max_val = val
            answers = [node]
        elif val == max_val:
            answers.append(node)    
    print(*answers)

solution1()

그런데, 메모리 초과가 뜬다. pyp3로 돌려 보라는 말이 있어 시도해 봤으나, 메모리 초과를 해결 할 수 없었다. 그래서 한 컴퓨터를 해킹했을 때 해킹되는 다른 컴퓨터들은 다시 그 컴퓨터들을 시작점으로 개수를 확인할 필요가 없으므로, dfs 중 방문한 컴퓨터들은 다시 해킹하지 않아도 되도록 코드를 수정해서 시도해보았다.

def solution2():
    N, relations = getInput()
    check = set([i for i in range(1, N+1)])
    num_lists = [-1] + [0 for _ in range(N)] #해킹한 컴퓨터 개수를 저장하기 위함
    while check: #더이상 확인할 컴퓨터가 없으면 끝
        start = check.pop()
        visited = [0 for _ in range(N+1)]
        visited[start] = 1
        dfs(start, relations, visited)
        num_lists[start] = sum(visited)
        check -= set([i for i, v in enumerate(visited) if v==1])
        #dfs중 방문한 컴퓨터들에서는 다시 dfs를 수행하지 않도록함
    max_val = 0
    answers = []
    for idx, val in enumerate(num_lists):
        if val > max_val:
            max_val = val
            answers = [idx]
        elif val == max_val:
            answers.append(idx)
    print(*answers)

메모리 초과는 되지 않지만 정답이 아니더라. 반례를 정말 오래 생각해봤는데, a->b->c->a 인 그래프인 경우라면 어떤 곳에서 시작하더라도 모두 해킹할 수 있으므로 정답이 a b c여야 하는데 내 코드대로 하면 a만 나오더라. 그래서 순환인 경우는 순환한 모든 경로의 위치에 대해 num_lists안에 동일한 개수를 넣어주면 해결 할 수 있어 보였다. 그래서 시도해 봤는데, 실력의 한계로 생각보다 복잡해져서 실패.

구글링을 통해 BFS로는 모든 컴퓨터에 대해 탐색해도 통과할 수 있다고 하여 다음과 같이 코드를 작성했다. def solution1의 dfs만 bfs로 수정하면 된다.

from collections import deque

def bfs(start, relations, visited):
    q = deque()
    q.append(start)
    visited[start] = 1
    while q:
        start = q.popleft()
        for end in relations[start]:
            if not visited[end]:
                q.append(end)
                visited[end] = 1