[프로그래머스] 다단계 칫솔 판매 - Python (tree 문제)

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

프로그래머스

푸는 로직은 간단하다 

1) 주어진 입력값으로 tree를 구성한다. 

친절하게도 문제에서는 enroll과 referrel 이 같은 순서로 주어진다. enroll[i] 이 자식노드, referrel[i] 의 원소가 부모노드가 된다 ! 

 

2) seller를 start로 하여 루트 노드 (center)까지 탐색한다 

재귀함수를 만들어서 자기 자신을 기준으로 tree에서 부모노드를 꺼내 부모노드로 이동하는 식으로 루트노드까지 방문하도록 코드를 짰다. 

 

함정(?) 이랄까 주의할 점은 

테스트케이스는 다 맞췄는데, 문제에서 요구하는 

- 절사단위는 원이고 배분하는 금액이 1미만이면 자기자신이 갖는다 

라는 조건을 꼭 처리해 주어야 한다. 

문제를 잘읽자.. 그리고 소숫점 버림 함수는 math 모듈에 trunc() 가 있다

from collections import defaultdict
import math 

def traverse(now, profit, tree, result):
    if now == "center": # 더이상 부모노드가 없는 경우 종료 
        return 
    share = math.trunc(profit*0.1) # 주의할점 !! 절사단위는 원이므로 소수점 다 버려야함
    if share < 1: # 나누는 이익이 1 미만이면 자기자신이 다 갖는다 
        result[now] += profit
        return 
    
    result[now] += (profit - share) # 자기자신한테 수익 (0.9)
    parent_node = tree[now]
    # 자기자신의 10퍼센트 
    traverse(parent_node[0], share, tree, result)

    
def solution(enroll, referral, seller, amount):
    n = len(enroll) # 노드 개수 
    m = len(seller) # 판매 정보 
    result = defaultdict(int)
    tree = defaultdict(list)
    
    for i in range(n):
        tree[enroll[i]] = []
    for i in range(n):
        if referral[i] != '-':
            tree[enroll[i]].append(referral[i])
        else: # root 노드가 부모노드인 경우 
            tree[enroll[i]].append("center")
    
    # start node 부터 root 노드까지 순회 
    for i in range(m):
        traverse(seller[i], amount[i]*100, tree, result)
    
    
    # 정답 배열 만들기 
    answer = []
    for i in range(n):
        answer.append(round(result[enroll[i]]))
    return answer