leetcode HARD : Vertical Order Traversal of a Binary Tree

https://leetcode.com/problems/vertical-order-traversal-of-a-binary-tree/

Vertical Order Traversal of a Binary Tree - LeetCode

 

친구랑 스터디하면서 같이 풀었다.

하드 문제 정말 오랜만에 도전해보는거같다!!

이 문제같은 경우는 어쨌든 position기준으로 정렬을 해야되기때문에 시간복잡도는 아무리 작아도 O(N log N)일거같다.

친구는 처음에 position기준으로 2차원 배열을 만들어서 나이브하게 푸는 방법을 떠올렸는데 공간복잡도가 너무 커질거같아서 포기!

나는 처음에 dict에 position을 key로 넣고 값을 value로 넣는 방법을 생각했는데 그러면 정렬이 안된다.

그래서 priority queue를 쓰는걸로 방법을 바꿨는데 생각해보니 그냥 list를 써서 정렬해도 상관없을거같다.

 

class Solution:
    def verticalTraversal(self, root: Optional[TreeNode]) -> List[List[int]]:
        pq = PriorityQueue()
        treemap = {}
        
        def search(node, pos):
            pq.put((pos, node.val))
            if node.left:
                search(node.left, (pos[0]-1, pos[1]+1))
            if node.right:
                search(node.right, (pos[0]+1, pos[1]+1))
                
        search(root, (0,0))
        
        answer = []
        tmp = []
        prevpos = -10000
        
        while not pq.empty():
            pos, val = pq.get()
            if pos[0] != prevpos:
                prevpos = pos[0]
                if tmp:
                    answer.append(tmp[:])
                    tmp = []
            tmp.append(val)
        
        if tmp:
            answer.append(tmp)
            
        return answer

 

핵심은 굵은글씨친부분! x랑 y 위치를 바꿔놨다. y가 앞에 와야 그거 기준으로 먼저 정렬할 수 있으니까.

우선순위 큐에 넣어버리고 꺼내면서 y가 같은것끼리 묶었다.