[Python 3] BOJ 31423: 신촌 통폐합 계획

이번 SUAPC WINTER 2024에 나왔던 H번 문제, 신촌 통폐합 계획이였다.

(다음 번에 SUAPC에 대해서 게재해보겠습니다.)

 

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

31423번: 신촌 통폐합 계획

처음에 이 문제를 봤을 때, 직관적인 풀이인 단순한 DFS로 접근했다. 

 

<틀린 코드>

import sys
sys.setrecursionlimit(int(5*1e5+1))
input=sys.stdin.readline
n=int(input())

names=[0]+list(input().strip() for _ in range(n))
data=[list() for _ in range(n+1)]
answer=[]
def dfs(num):
    answer.append(num)
    for i in data[num]:
        dfs(i)
first=0
for i in range(n-1):
    a,b=map(int,input().split())
    data[a].append(b)
    if i==n-2:
        first=a
realanswer=""
dfs(first)
for i in range(n):
    print(names[answer[i]] ,end="")

 

이 코드로 pypy3로 제출하게 되면 메모리 초과를, python 3로 제출하게 되면 시간 초과를 받게 된다.

원인은 간단하다. 최대 50만개의 학교 이름이 나올 수 있으니, python의 기본 recursion limit인 1000개를 넘어서 50만개로 바꿔줘야 하는데, 미리 메모리를 선점하고 나서 코드를 돌리는 방식이기 때문에 결국 메모리 초과가 나게 된다. (근데 python이 아니라 c++로 하면 해당 코드가 통과한다. 즉, 파이썬 한정 언어 디메릿이 있다고 할 수 있다.)

 

해당 코드를 linear time에 해결 할 수 있도록 변경해보겠다.

import sys
input=sys.stdin.readline
n=int(input())
name=[0]+list(input().strip() for _ in range(n))
lastindex=[i for i in range(n+1)]
nextindex=[i for i in range(n+1)]
for _ in range(n-1):
    a,b=map(int,input().split())
    
    nextindex[lastindex[a]]=b
    lastindex[a]=lastindex[b]
    
for i in range(n):
    print(name[a],end="")
    a=nextindex[a]

 

 

dfs가 아닌 linked list를 implement해보았다. 

lastindex는 각 인덱스의 마지막 대학, nextindex는 바로 다음에 오는 인덱스의 대학이다.

a와 b라는 대학이 for문에서 주어진다.

a는 a,a',a''가, b는 b,b'로 이루어져있다고 하자.

a,a',a'',b,b' 순서대로 이어져야한다.

당연히 a'', 즉 a의 마지막 인덱스, 즉 lastindex[a]의 다음 index는 b여야한다.

그 이후, a의 마지막 인덱스, 즉 lastindex[a]를 b의 마지막 인덱스 (b')로 설정해야 하기에, lastindex[a]=lastindex[b]로 설정한다.

 

쉬우면서도 생각할게 많았던 문제였다.