[프로그래머스] 1,2,3 떨어뜨리기, python

https://school.programmers.co.kr/learn/courses/30/lessons/150364

프로그래머스

 

공을 굴렸을 때 도달하는 리프노드의 순서를 배열로 저장한다.

처음 떨어뜨리면 노드 4에 도착하고, 두 번째 떨어뜨리면 바뀐 루트를 따라서 노드 7에 도착한다. 이런 식으로 dfs를 탐색할 때 그래프의 간선위치를 갱신하면서 도착한 값을 배열에 저장하면

4,7,6,8,4,10,7... 이 저장된다.

 

위에서 도착한 노드 배열을 통해 각 노드에 몇번 도착했는지를 저장하는 배열을 생성한다.

4, 7, 6, 8, 4라면 4에 2번, 6에 1번, 7에 1번, 8에 1번을 저장한다.

 

리프노드에 도달한 횟수가 target에 유효범위에 안에 드는 최소 턴을 구한다.

도달 횟수가 n이라면 k과 3k 범위 내를 만족해야 한다.

위에서 4번 노드의 도달 횟수가 2번이라면 target은 최소 2, 최대 6 이어야 한다.

턴을 하나씩 늘려가며 모든 노드에 대해 위 조건을 만족하는 최소턴을 구한다.

 

target값을 리프노드에 도달한 횟수로 분배 후 오름차순정렬한 배열을 생성한다.

target이 5고 도달한 횟수가 3이라면 5를 3개로 분배해야 한다.

5를 3개로 분배하는 법은 1,1,3  / 1,2,2 가 있는데 여기서 공의 순서를 가장 적은 수가 앞에 와야 한다.

분배하려는 수(예시에서는 3개) 만큼 1로 분배한 후 뒤쪽부터 남은 값을 채워나가면 최솟값이 앞에 오게 된다.

위의 예시로 보면 5를 1,1,1로 분배 후에 5에서 3개를 뺀 2개를 뒤쪽부터 분배하면 1, 1, 3이 된다. 

여기서 저장된 값은 해당 리프노드에 도착한 공 순서가 된다.

 

마지막으로 맨 처음에 저장한 도착한 노드 순서를 저장한 배열에서 순서대로 돌면서 리프 노드를 방문 시 위에서 구한 공 순서 배열에서 하나씩 꺼낸다.

여기서는 꺼내는 것 대신 재방문 시 인덱스를 1개씩 늘리는 것으로 구현했다.

 

import sys
sys.setrecursionlimit(10**6)

def solution(edges, target):
    answer = []
    
    n = len(target)
    graph = [[] for _ in range(n)]
    edgeIndex = [0]*n # 해당 노드가 가리키는 자식노드 저장
    
    for x,y in edges:
        graph[x-1].append(y-1)
    
    for x in graph:
        x.sort()
        
    leafIndex = [] # 공을 떨어뜨렸을때 도달하는 리프노드들의 순서를 저장할 배열
    
    def dfs(cur):
        if len(graph[cur]) == 0:
            leafIndex.append(cur)
            return
    
        next = graph[cur][edgeIndex[cur]]
        # 간선을 한칸 옮김
        edgeIndex[cur] = (edgeIndex[cur]+1)%(len(graph[cur]))
        dfs(next)
    
    # 떨어뜨릴 공의 총 합만큼 리프노드 배열을 생성
    for _ in range(sum(target)):
        dfs(0)
    
    check = [False]*n
    cnt = [0]*n # 각 노드에 공이 몇번 도달했는지 저장
    turn = 0 # 조건을 만족할때 까지 걸리는 턴수
    for x in leafIndex:
        cnt[x] += 1
        # 리프노드에 target값이 유효범위 내에 있다면 True
        if cnt[x] <= target[x] <= cnt[x]*3:
            check[x] = True

        # 리프노드에 target값이 유효범위 밖이면 False
        else: 
            check[x] = False
        
        # target에 값이 있는 개수와 check(유효범위)가 True인 개수가 같으면 break
        if sum([1 for x in check if x]) == sum([1 for x in target if x!=0]):
            break
        turn += 1

        # 유효하지 않다면 -1
    else:
        return [-1]
        
    
    # 해당노드에 합을 노드에 공이 도달한 횟수로 분배함
    # 모두 1을 할당하고 뒤에서부터 채워나감 
    def devide(num, cnt):
        result = [1]*cnt
        a = num-cnt
        for i in range(cnt):
            if a==0:
                break
            if a>=2:
                result[i] += 2
                a -= 2
            elif a>=1:
                result[i] += 1
                a -= 1
        result.sort()
        return result  
    
    
    newcnt = cnt[:]

    for x in range(turn+1):
        i = leafIndex[x] # 해당 턴에서 굴렸을때의 리프노드
        if target[i] == 0:
            continue
        
        # target을 리프노드도달 수로 분할한 배열
        # 해당 노드에 도달한 순서대로 저장
        t = devide(target[i],cnt[i])
        
        # t배열에서 하나씩 넣음, 해당노드에 재방문시 순서대로 넣기위함
        answer.append(t[len(t)-newcnt[i]])
        newcnt[i] -= 1
    
    return answer