979.Distribute Coins in Binary Tree

In [4]:

import sys
from binarytree import Node
sys.path.append("/home/swyoo/algorithm/")
from utils.verbose import buildtree

979. Distribute Coins in Binary Tree

Idea

우리의 목표는 결국 모든 노드마다 코인이 하나씩 있도록 분배하는데 코인을 몇번 넘겨주냐를 구하는 것. 노드마다 남게되는 코인수를 return 하자.
이 과정에서 ans를 업데이트 하면 된다.
결국 모든 코인이 동일하게 1개씩 분배되기 때문에 (분배되지 않는 경우는 없음) left 트리에서 right 트리에서 각각 필요한 수, 또는 남는 수만큼씩을 ans에 업데이트하면 된다. (이 부분이 생각하기 까다롭다. left, mid, right 에 대해 경우의 수를 나눠 생각해보면
ans += (abs(left) + abs(right)) 만 해주면 됨을 알 수 있다.)

In [5]:

class Solution:
    def distributeCoins(self, root: Node) -> int:
        ans = 0
        def dfs(cur: Node):
            """ return remainder coins, if deficient, negative value.
            update ans value through search process. """
            nonlocal ans
            if not cur:
                return 0
            left = dfs(cur.left)
            right = dfs(cur.right)
            mid = cur.value - 1
            rmd = left + right + mid
            ans += (abs(left) + abs(right))
            return rmd
        dfs(root)
        return ans
    
sol = Solution()

In [10]:

null = None
A = [1,0,0,null,3]
# A = [1,0,2]
root = buildtree(A)
root.pprint()
print("ans:", sol.distributeCoins(root))

  __1
 /   \
0     0
 \
  3

ans: 4