p_시험장나누기_81305

문제

유니온파인드하면서 합치면될듯?

답

내답

import heapq

def solution(k, num, links):
    N = len(num) # 노드의 개수
    parent = [i for i in range(N)]  # 직계 부모
    weight = [] 
    for i in range(N):
        weight.append([-num[i],i])
    for i in range(N):
        l,r = links[i]
        increase= 0# 음수임 
        if l > -1:
            parent[l] = i
            increase += weight[l][0]
        if r > -1:
            parent[r] = i
            increase += weight[r][0]
        weight[i][0] += increase
        update(parent,weight,increase,i)
        # print(i,"------")
        # print(parent,weight,increase)
    
    heapq.heapify(weight)
    # heap 사용하고 최소힙이니깐 - 이용함
    # k만큼 for 문
    # 앞에꺼하나 뺌
    # 자식이없음. > 걍그거 리턴하시면됨.
    # 자식보다 부모가 항상큼 > 고로 pop하면서 아니면 다시넣어주고 이런식으로하면댐
    # 자식이잇음 > 자식이 하나면 그거랑 끊고 자식이 2명이면 둘중 큰거랑 끊음> 나랑 직계부모 업데이트, 부모정보도 업데이트+링크정보도 업뎃
    # print('w',weight)
    for _ in range(k-1):
        w,i = heapq.heappop(weight)
        l,r = links[i]
        if parent[l] != i:
            # 이미 끊어 졋던거임
            l = -1
            links[i][0]= -1
        if parent[r] != i:
            r = -1
            links[i][1]= -1
        if l<0 and r<0:
            return -w
        else:
            targetW = None 
            depth = 0
            remain = []
            while targetW is None:
                # print(depth)
                if depth >N:
                    break
                depth +=1
                nw,ni = heapq.heappop(weight)
                if ni <mark> l or ni </mark> r:
                    targetW = [nw,ni]
                else:
                    remain.append([nw,ni])
                    
            for rem in remain:
                heapq.heappush(weight,rem)
            # print(targetW)

            # 부모랑 연결끊기
            if targetW is None:
                return -w
            tw,ti=targetW
            w -= tw
            parent[ti] = ti
            heapq.heappush(weight,[w,i])
            heapq.heappush(weight,targetW)
        #     print(weight)
        #     print(parent)
        # print("-")
        
    return -weight[0][0]

def update(parent,weight,value,idx):
    i = idx
    while parent[i] != i:
        i = parent[i]
        weight[i][0] += value

실제 답
https://loosie.tistory.com/342