32224:물탱크 알바(easy) [백준 문제풀이][Python]

해설지에서 미쳐 다 담지 못한 디테일한 풀이를 정리해서 올려보았습니다. 혹시라도 읽어보고 이해가 되지 않는 부분이나 다른 의견은 댓글에다가 남겨주시길...

 

문제: https://www.acmicpc.net/problem/32224

난이도: 골드 상위 ~플레 하위

Code Master 2024에 내가 직접 출제한 문제이다. DFS를 사용하여 풀면 된다.

뭔 문제야 이게?

혹시라도 문제를 이해하지 못 한 사람들을 위해 예제를 통해 설명해주겠다.

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

예제 1의 도식화

예제입력 1을 tree 구조로 나타내면 이러한 형태일 것이다.

이 상황에서 6의 물을 어느 물탱크부터 넣어야 가장 많은 물탱크를 채울 수 있을까? 중요한 조건은 현태 물탱크가 꽉 찬다면, 부모 물탱크로 물이 넘쳐 올라간다는 것이다. 우선 1번 물탱크에 물을 넣어보자.

1번 물탱크의 물을 넣었을 경우

2,3의 물탱크로 3,3의 물이 흘러 들어가서 아래 물탱크부터 차기 시작하여 아래 있는 4개의 물탱크가 찬다.

이번에는 2번 물탱크에다가 6의 물을 넣어보자.

2번 물탱크에 물을 넣었을 경우

2,4,5 번 물탱크가 먼저 채워지고, 남은 2의 물이 1번 물탱크로 올라가 반대쪽 물탱크로 흘러 들어간다.

그렇게 6,7번 물탱크가 채워지고 최종적으로 5개의 물탱크가 채워진다. 결국 최대 해는 5가 된다.

 

우선 관찰부터

우리가 시간복잡도를 줄이기 위해 해줘야 할 질문은 과연 모든 물탱크를 체크해주어야 하는가 이다.

다음과 같은 상황을 관찰해보자.

3 7
4 4 3 2
1 1 2

 

2번 / 4번 물탱크에 물을 넣었을 경우

위 그림은 2번, 4번 물탱크에다 물을 넣었을 때의 그림이다. 이를 통해 알 수 있는건 물이 꽉 찬다면 부모 물탱크로 올라가므로 가득 차는 하위 물탱크에 대해서는 탐색을 해주지 않아도 된다는 것이다. 그렇다면 어느 부분에서 탐색을 해주어야 할까? 1번 ,3번 물탱크를 채우는 경우를 보자.

1번 / 3번 물탱크에 물을 넣었을 경우

위 그림들에서 알 수 있듯이 자식 물탱크가 2개여서 부모 물탱크를 통해 반대쪽 물탱크로 물이 넘어가는 분기점에서 차이점이 생긴다는 것을 알 수 있다.

따라서 각 물탱크에 물을 넣었을때 넘치는지 안 넘치는지를 확인한 후, 만약 넘치지 않는다면 그 물탱크과 자식 물탱크에 대해서 체크를 해주면 된다.

 

이 물탱크에 물을 넣는다면?

 물론 탐색을 할때마다 하위 물탱크까지 찍고 올라오는 행위를 반복 할 수 없다. 따라서 우리는 DFS를 이용하여 각 물탱크의 x개의 물탱크를 채울려면 몇의 물을 채워야 하는지를 하위노드에서부터 체크해주며 올라갈 것이다.

이제 각 물탱크 마다 리스트 tank[n]를 할당해 줄 것이다. tank[n][x]은 n번 물탱크에다가 x개의 물탱크를 채울 때 필요한 물의 양이다.

자식 물탱크가 0개일때

우선 자식이 없는 리프 물탱크의 경우이다. 0개를 채울때는 0의 물이 필요하고, 자기 자신을 채우는 경우를 뒤에 추가해주면 된다.

자식 물탱크가 1개일때

다음은 자식이 한 개인 물탱크의 경우이다. 자기 자신을 채울려면 무조건 자식 물탱크들을 전부 채워야 하므로, 자신의 자식 물탱크의 리스트를 가져온 후 마지막 요소에다가 자신의 용량을 더하여 추가해주면 된다.

자식 물탱크가 2개일때

마지막으로 자식이 두 개인 물탱크이다(이진 트리이므로 최대가 2개). 이 경우는 채우는 과정을 생각하며 상황을 나눠줘야 한다.

1. 자식 물탱크 2개가 완전히 안 차서 들어오는 물이 정확히 반으로 나누어 들어가는 경우

2. 자식 물탱크 1개가 완전히 찬 후 남은 물탱크에 들어오는 물이 온전히 들어가는 경우

 

1번 상황을 생각해보면 자식 물탱크에서 차례대로 작은 값을 X2 하여 리스트에 넣어주면 된다. 그러다 한 쪽 물탱크가 완전히 다 차는 순간, 2번 상황이 된다. 근데 2번 상황은 위에서 얘기한 자식 노드가 한 개인 상황과 같기 때문에 남은 자식 물탱크의 요소들을 부모 물탱크의 마지막 요소에 더하여 넣어주면 된다.

 

이런 과정을 거치면 한 번의 탐색만으로도 리스트를 참조하여 얼마나 많은 물탱크가 차는지 알 수 있다.

 

물이 꽉 차지 않는다면

위 과정을 통해 리스트를 만들어가다가 리스트의 값이 주어진 물의 값인 m을 넘어가거나, 루트 물탱크에 도달하는 경우가 있을 것이다. 이 말은 모든 물을 넣어도 현재 물탱크가 꽉 차지 않는다는 말이므로, 부모 물탱크와 자식 물탱크에 대해 탐색을 해주어야 한다.

부모 물탱크에 물을 넣는 경우

부모 물탱크에 물을 넣는 경우, 이진 탐색을 사용하여 최대 물탱크 개수를 찾아주면 된다. 위 그림과 같이 답은 1개가 된다.

자식 물탱크에 물을 넣는 경우

자식 물탱크에 넣는 경우, 2번의 이진 탐색을 해주면 된다. 한 쪽에서 이진탐색을 통해 물이 넘어 갈 수 있는지 판별해주고, 넘어간다면 남은 물을 다른 자식 물탱크에서 이진 탐색을 이용해 몇 개를 채울 수 있는지 확인하여 두 값을 더해주면 된다. 위 경우는 왼쪽에서 5의 물을 이용하여 2개, 오른쪽에서 2의 물을 이용하여 0개를 채웠으므로 답은 2개가 된다.

 

코드로 옮기기

 위 과정을 코드로 옮기면 다음과 같다.(물론 내 코드가 최선은 아닐 것이다.)

import sys
input = sys.stdin.readline
from bisect import bisect_right
sys.setrecursionlimit(10**8)

n,m = map(int,input().split())
a = list(map(int,input().split()))
parent = list(map(int,input().split()))

p = [[] for _ in range(n+1)]
for i in range(0,n-1):
    p[parent[i]-1].append(i+1)
root = 0
answer=0
tank = [[] for _ in range(n+1)]


def dp(x):
    global answer
    if(len(p[x])==0):#자식 물탱크가 0개라면
        tank[x].append(0)
        tank[x].append(a[x])
        
        
    elif(len(p[x])==1):#자식 물탱크가 1개라면
        y1 = p[x][0]
        dp(y1)
        for i in tank[y1]:tank[x].append(i)
        tank[x].append(tank[x][-1]+a[x])
        
        
    elif(len(p[x])==2):#자식 물탱크가 2개라면
        y1,y2 = p[x][0],p[x][1]
        dp(y1)
        dp(y2)
        tank[x].append(0)
        i,j=1,1
        '''2번 상황이 될 때까지 *2 하며 추가해주기'''
        while(i<len(tank[y1]) and j<len(tank[y2])):
            if(tank[y1][i]==tank[y2][j]):
                tank[x].append(tank[y2][j]*2)
                tank[x].append(tank[y1][i]*2)
                i+=1
                j+=1
            elif(tank[y1][i]>tank[y2][j]):
                tank[x].append(tank[y2][j]*2)
                j+=1
            elif(tank[y1][i]<tank[y2][j]):
                tank[x].append(tank[y1][i]*2)
                i+=1
        
        '''남은 요소들 처리해주기'''
        if(j<len(tank[y2])):
            tank[x].append(tank[x][-1] // 2 + tank[y2][j])
            for k in range(j,len(tank[y2])-1):tank[x].append(tank[x][-1]+ (tank[y2][k+1]-tank[y2][k]))
        elif(i<len(tank[y1])):
            tank[x].append(tank[x][-1] // 2 + tank[y1][i])
            for k in range(i,len(tank[y1])-1):tank[x].append(tank[x][-1]+ (tank[y1][k+1]-tank[y1][k]))
        tank[x].append(tank[x][-1]+a[x])
        
        
    if(tank[x][-1]>=m or x==root):#m을 넘어가거나 root일때 체크해주기
        pb=bisect_right(tank[x],m)-1
        b1,b2=0,0
        if(len(p[x])>=1):b1 = bisect_right(tank[y1],m)-1
        if(len(p[x])>=2):
            b2 = bisect_right(tank[y2],m)-1
            if(tank[y2][-1]<m):b2 += (bisect_right(tank[y1],m-tank[y2][b2]) - 1)
            if(tank[y1][-1]<m):b1 += (bisect_right(tank[y2],m-tank[y1][b1]) - 1)
        b = max([pb,b1,b2])
        if(b > answer):answer = b
dp(root)

print(answer)

 

마치며

최대한 읽어보기 쉽게 해설을 적어보긴 했는데, 난해한 부분이 있진 않는지 모르겠습니다. 혹시라도 이해가 되지 않거나 모르겠는 부분이 있다면 댓글로 남겨주시길 바랍니다. 마지막으로 오프라인, 오픈 콘테스트에 참여해주신 모든 분들과 이 문제를 푼 모든 분들께 감사드립니다!