LeetCode [ 637 | Python ] Average of Levels in Binary Tree

문제

https://leetcode.com/problems/average-of-levels-in-binary-tree/description/?envType=study-plan-v2&envId=top-interview-150

주어진 이진 트리에서 각 계층마다의 평균값을 반환하는 문제.

풀이

시도 1

실행시간 : 43ms 메모리 : 18MB

계층적으로 로직을 처리해야하기 때문에 BFS로 순회했다.
값을 하나씩 딕셔너리 dic에 계층별로 저장한 후에, 마지막으로 이를 더한데서 그 길이로 나누어 리스트를 반환했다.

정말 아무생각 없이 일단 적기시작하면 이런 코드도 나온다.
실행시간도 실행시간이겠지만, 메모리 측면에서도 하나도 고민안한 코드라고 생각된다.

첫번째 목표로 먼저 마지막 줄 sum 함수를 제거해보는 것을 목표로 삼았다.

from collections import deque
class Solution(object):
    def averageOfLevels(self, root):
        depth = 0
        dic = {}

        def bfs(start):
            q = deque()

            q.append([start, 0])
            
            while q:
                x, d = q.popleft()

                if d in dic:
                    dic[d].append(x.val/1.0)
                else:
                    dic[d] = [x.val/1.0]
                
                if x.right:
                    q.append([x.right, d + 1])
                if x.left:
                    q.append([x.left, d + 1])
                
        bfs(root)
        return list(map(lambda x : sum(x[1])/len(x[1]), dic.items()))

개선 1

실행시간 : 38ms 메모리 : 17.7MB

딕셔너리에 원소를 추가하는 것이 아닌 그 값을 더하는 형식으로 개선했다.
또한 그 대신 딕셔너리에 해당 깊이에서의 원소 개수를 따로 저장했고, 이 덕분에 sum과 len을 제거할 수 있었다.

하지만 아직도 개선점이 남아있다.
이전 문제도 그래왔지만, 깊이를 따로 저장하고 있다는 점이다. 다음은 이를 제거하는 것을 목표로 삼았다.

from collections import deque
class Solution(object):
    def averageOfLevels(self, root):

        depth = 0
        dic = {}

        def bfs(start):
            q = deque()

            q.append([start, 0])
            
            while q:
                x, d = q.popleft()

                if d in dic:
                    dic[d][0] += x.val
                    dic[d][1] += 1
                else:
                    dic[d] = [x.val/1.0, 1]
                
                if x.right:
                    q.append([x.right, d + 1])
                if x.left:
                    q.append([x.left, d + 1])
                

        bfs(root)
        
        return list(map(lambda x : x[0]/x[1], dic.values()))
        

개선 2

실행시간 : 18ms 메모리 : 17.7MB

코드를 유심히 보다보니 잊었던 사실이 q를 한번 다 채웠을 때는 모두 같은 깊이에 있는 노드들이라는 사실이었다.
이름부터도 너비 우선 탐색인데도 잊었다.
이렇게 된다면 본래 트리의 크기를 알 수 없어 사용했던 딕셔너리도 사용할 필요 없다.

그에 따라 한 원소를 통해 left, right에 존재하는 모든 노드들을 큐에 저장하고 그 크기와 총합을 나눠 res 리스트에 저장한다.
깊이별로 평균값을 담을 수 있는 것이다.

from collections import deque
class Solution(object):
    def averageOfLevels(self, root):
        def bfs(start):
            res = []
            q = deque()
            q.append(start)
            
            while q:
                size = len(q)
                total = 0.0

                for i in range(size):
                    x = q.popleft()
                    total += x.val

                    if x.left:
                        q.append(x.left)
                    if x.right:
                        q.append(x.right)

                res.append(total/size)
 
            return res
 
        return bfs(root)
        

생각 & 정리

때로 가장 기본적인걸 잊어서 비효율을 빚거나 고생을 크게 하는 경우가 있다.

해결할 수 있는 부분인데 ..

라는 생각이 든다면 놓지말고 끝까지 보자는 생각을 다잡을 수 있었다.