2021 KAKAO 채용연계형 인턴쉽 시험장 나누기

https://programmers.co.kr/learn/courses/30/lessons/81305#

코딩테스트 연습 - 시험장 나누기

 

Naive한 방법에서 현재 풀이까지 생각하는데 많은 시간(2주)이 걸렸다.

Naive한 방법은 간선을 끊어 모든 경우의 수를 탐색하는 방법이다. 이 방법은 정확성 Test에서는 합격을 받을 수 있겠지만, 효율성에서 받지 못할 것 같았다.

다음으로 생각한 방법은 DP를 이용하는 것이였다. 하지만, DP를 쓰더라도 간선 처리, dp의 index를 무엇으로 잡을거냐에 대한 생각이 나지 않았다.

마지막으로 생각한 방법은 parametric search였다. "k개로 끊었다고 했을때, 분리된 시험장의 최대값이 L이하인가?" 라고 문제를 재정의하고 생각을 했지만, 이 방법도 결국 "끊는 간선을 어떻게 선택하고, 그 후 간선들도 어떻게 처리해야할까?" 라는 생각이 들어 적용하지 못했다.

하지만 며칠 전에 "끊는 간선을 tracking하는 것이 진짜 중요한가?" 에 대해 생각이 들었다. parametric search를 사용했을때 중요한 것은 해당 조건을 만족하는 것을 확인하는 것이다. 

그래서 parametric search의 가정을 바꿔서 "분리된 시험장의 최대값이 L이하일 때, k개로 끊을 수 있나?" 로 생각하니, button-up으로 끊어 올라가면서 L이상의 subtree들이 생기면 버리는 방식으로 하면 될 것 같다고 생각했다.

또한, "꼭 k개로 끊는것만 확인해야할까? 라는 생각이 들었다." k개 이하로 끊을 수 있다면 L을 줄여서 더 tight하게 보면 되고, k개로 끊을 수 있는 방법이 없다면 L을 늘려서 rough하게 보면 된다.

이런 방식으로 생각하니, DP식을 세울 수 있게 되었다. 이 문제에서 처음으로 DP식을 2개를 정의하게 되었다.

cnt[ i ] = 현재 시험장을 제외한 i node의 subtree의 시험장 수
total[ i ] = i node를 포함한 시험장의 총 응시인원수 

경우의 수는 다음과 같다.

0. 모든 경우에 다음이 성립한다.
cnt [ i ] = cnt [left child index] + cnt [right child index] 

1. num [ i ] + total [ left child index] + total [right child index] < L
이 경우는 left child 의 시험장, right child의 시험장, 현재 node를 한 시험장으로 구성할 수 있다.
total [ i ]  = num [ i ]  + total [ left child index] + total [right child index] 

2. num [ i ] + total [ left child index ] < L
이 경우는 left child의 시험장, 현재 node를 한 시험장으로 구성할 수 있다.
또한, right가 잘려나가기 때문에 cnt [ i ] 에 1을 추가해준다.
total [ i ]  = num [ i ]  + total [ left child index]
cnt [ i ] += 1

3. num [ i ] + total [ right child index ] < L
case 2의 반전 버전.

4. num [ i ] < L
이 경우는 현재 node로만 시험장을 구축할 수 있다.
또한, left, right가 잘려나가기 때문에 cnt[ i ] 에 2를 추가해준다.
total [ i ]  = num [ i ] 
cnt [ i ] += 2

5. num [i] > L
이 경우에는 시험장의 최대 응시생을 L이하로 구축할 수 없다. 그래서 바로 return 해준다.

(이 외에도 left child /right child 가 없는 경우는 처리를 따로 해준다.)

parametric search는 cnt [ root ] + 1 (+1은 root node의 시험장 수) 로 값 범위를 설정한다.

이렇게 하면 틀린다고 나온다...
왜 틀릴까 고민했는데, case 2, case 3을 좀더 자세히 보니, 이 과정에서 선택이 필요한 부분이 있었다. 반례는 우측을 선택해도, 좌측을 선택해도 만족하는 경우였다. 이 경우는 subtree의 시험장 수가 적은 것을 선택하는 것이 유리하다.

이를 적용해서 문제를 풀었다.

코드는 다음과 같다.

import sys
sys.setrecursionlimit(100000)
INF = 10000000000

def solution(k, num, links):
    answer = 0
    n = len(num)
    total = [0] * n
    cnt   = [0] * n
    root = -1
    indgree = [0] * n
    for link in links:
        for l in link:
            if l != -1:
                indgree[l] += 1
    for idx in range(n):
        if not indgree[idx]:
            root = idx
            break
            
    
    
    def dfs(now, target):

        for link in links[now]:
            if link != -1:
                flag = dfs(link, target)
                if not flag: return False
                
            else: continue
        
        link = links[now]
        
        left = INF if link[0] == -1 else total[link[0]]
        right = INF if link[1] == -1 else total[link[1]]
        
        #print(now, num[now], left, right, target, link)
        cnt[now] = (0 if link[0] == -1 else cnt[link[0]]) + \
                    (0 if link[1] == -1 else cnt[link[1]])
        if num[now] + left + right <= target:
            #print("case 1")
            total[now] = left + right + num[now]
        elif num[now] + left <= target or num[now] + right <= target:
            if link[0] == -1:
                #print("case 2-1 right only")
                total[now] = right + num[now]
            elif link[1] == -1:
                #print("case 2-2 left only")
                total[now] = left + num[now]
            else:
                #print("case 2-3 select minimum")
                total[now] = num[now] + min(right, left)
                cnt[now] += 1
            
        elif num[now] <= target:
            #print("case 4")
            total[now] = num[now]
            
            if link[0] != -1 : cnt[now] += 1
            if link[1] != -1 : cnt[now] += 1
            
        else:
            #print("case 5")
            return False
        
        return True
    
    
    start = 0
    end   = sum(num)
    answer = sum(num)
    while start <= end:
        mid = (start + end) // 2
        #print(start, mid, end)
        flag = dfs(root, mid)
        '''
        [start, mid], [mid+1, end]
        '''
        
        if not flag or cnt[root]+1 > k:
            #print("increase")
            start = mid + 1
        else:
            #print("decrease")
            end = mid
            answer = min(answer, max(total))
        if start == end:
            break
    
    return answer