2265. Count Nodes Equal to Average of Subtree

문제링크: https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree/description/

숫자가 들어있는 트리가 주어진다. 트리의 각 노드 중 자신의 값과 부분트리의 값의 평균이 같은 노드들의 개수를 구하는 문제다. 이때 평균은 내림 해 같은지 확인한다.

Solution

Postorder dfs

어떤 노드에서 서브트리의 평균을 구하기 위해서는 총 합과 노드의 개수가 필요하다. 각 노드에서의 합과 노드 개수는 탐색으로 구할 수 있다. 합과 노드 개수는 아래서부터 연산되어 전파할 수 있기 때문에 bottom-up 방식으로 접근하는게 효율적이다. 후위 순회를 통해 트리에서 bottom-up 방식으로 구현해보았다.

Algorithm

• 조건에 맞는 노드 수를 연산할 result를 자유변수로 설정한다.
• dfs함수를 통해 노드를 순회한다.
• 노드가 null인 경우 [0, 0] (totalSum:0, totalCount:0)을 반환한다.
• 후위순회를 위해 좌우 노드를 재귀적으로 우선 탐색한다.
• 후위순회로 전파받은 데이터를 받아온다.
• 현재 노드에서 totalSum과 totalCount를 전파된 데이터로 얻는다.
• 조건에 부합하는지(평균 내림값과 노드값이 일치하는지) 확인하고 result에 1을 더한다.
• 이후 탐색에 필요하도록 totalCount와 totalSum을 전파해준다.
• 탐색이 끝난 후 result를 반환한다.

var averageOfSubtree = function(root) {
  // 노드가 자신이 부모인 subtree의 평균이 자신 값과 같은 경우를 센다. 
  // Postorder traversal
  // Total sum && node count 2가지 정보로 평균을 구할 수 있다.
  // 이는 bottom-up 즉 postorder로 계산하면 아주 간단하다.

  let result = 0;

  const dfs = (node) => { 
  // 후위순회를 하면서 subtree의 평균이 자신의 노드값과 같을 때 result 카운트를 +1 해주는 방식
    if (node === null) return [0, 0];
    const [leftSum, leftCount] = dfs(node.left);
    const [rightSum, rightCount] = dfs(node.right);

    // 평균이 자신의 값과 같을 때 result++
    const totalSum = leftSum + rightSum + node.val;
    const totalCount = leftCount + rightCount + 1;
    if (Math.floor(totalSum / totalCount) === node.val) result++;

    // 값을 넘겨줘서 후위순회시에 연산을 할 수 있도록
    return [totalSum, totalCount];
  }

  dfs(root); // 전체 순회
  return result;
};

시간복잡도는 각 노드를 한번씩만 탐색하기 때문에 O(n)이다. 공간복잡도는 깊이 탐색을 위해 재귀적으로 생긴 스택 영역으로 인해 O(height)가 된다.