[LeetCode] 124. Binary Tree Maximum Path Sum (Hard)

이 문제를 풀기 위해서는 후위 순회(postorder traversal)와 트리 DP에 대해서 알아야 한다.

각 노드에서 자식들이 만들어 둔 최적의 값을 받아 현재 노드의 값을 계산한다. 다만 이 문제에서는 부모에게 넘길 값과 전체 정답 후보가 서로 다르다는 점이 핵심이다.

0. 필요한 개념 먼저 정리하기

0-1. 경로와 후위 순회

트리의 경로는 인접한 노드를 간선으로 이어 만든 노드의 나열이다. 같은 노드를 두 번 지날 수 없고, 루트를 반드시 포함할 필요도 없다. 따라서 어떤 노드에서 왼쪽 자식으로 내려갔다가 다시 그 노드를 거쳐 오른쪽 자식으로 내려가는 형태도 하나의 유효한 경로다.

현재 노드에서 최선의 값을 계산하려면 왼쪽과 오른쪽 서브트리의 결과가 먼저 필요하다. 자식을 먼저 처리하고 부모를 나중에 처리하는 순서가 후위 순회다.

0-2. 트리 DP와 기여값(gain)

각 노드 node에 대해 다음 값을 저장한다고 생각해보자.

gain(node) = node에서 시작해 아래쪽으로 한 방향만 이어지는 경로의 최대 합

이 값은 부모의 경로에 이어 붙일 수 있어야 한다. 부모에서 내려온 경로가 현재 노드에서 왼쪽과 오른쪽으로 동시에 갈라지면 하나의 단순 경로가 아니므로, 두 자식 중 더 큰 한쪽만 선택해야 한다.

반면 현재 노드를 경로의 가장 높은 꼭짓점으로 삼는다면 왼쪽과 오른쪽을 모두 연결할 수 있다. 이 값은 부모에게 넘기지는 못하지만 전체 정답의 후보가 된다.

1. 접근 : 문제를 단순화 하기

주어진 이진 트리에서 비어 있지 않은 경로 하나를 골라 노드 값의 합을 최대화해야 한다. 노드 수는 최대 30,000개이고 값에는 음수도 포함된다.

먼저 문제를 조금 단순화해서 생각해보자. 어떤 노드에서 부모 쪽으로 연장할 수 있는 경로는 세 가지뿐이다.

1. 현재 노드만 사용한다.
2. 현재 노드와 왼쪽의 최적 기여값을 사용한다.
3. 현재 노드와 오른쪽의 최적 기여값을 사용한다.

왼쪽과 오른쪽을 모두 붙인 경로는 현재 노드에서 갈라지므로 부모에게 다시 연결할 수 없다. 대신 그 형태 자체는 완성된 경로로서 전역 최댓값을 갱신할 수 있다.

모든 경로를 직접 열거하면 경우의 수가 너무 많다. 각 노드에서 필요한 두 값만 한 번씩 계산하면 전체를 O(n)에 해결할 수 있다.

2. 핵심 아이디어

왼쪽 자식이 주는 기여값을 leftGain, 오른쪽 자식이 주는 기여값을 rightGain이라고 하자. 음수 기여값을 붙이면 합이 작아지므로 0과 비교해 버린다.

const leftGain = Math.max(0, gain.get(node.left) ?? 0);
const rightGain = Math.max(0, gain.get(node.right) ?? 0);

현재 노드를 꼭짓점으로 하는 완성된 경로는 다음과 같다.

leftGain + node.val + rightGain

이 값으로 전체 정답을 갱신한다. 부모에게는 양쪽 중 더 큰 한쪽만 이어서 반환한다.

node.val + max(leftGain, rightGain)

여기서 중요한 포인트는 전역 정답을 0이 아니라 -Infinity로 시작해야 한다는 것이다. 모든 노드 값이 음수라면 빈 경로는 허용되지 않으므로, 가장 덜 작은 노드 하나가 정답이어야 한다.

재귀 DFS로도 같은 점화식을 구현할 수 있지만, JavaScript에서는 높이가 30,000인 편향 트리가 호출 스택 한도를 넘길 수 있다. 따라서 [node, visited]를 담는 명시적 스택으로 후위 순회를 구현한다.

3. 단계별로 구현하기

3-1. 후위 순서 만들기

처음 만난 노드는 visited = true 상태로 다시 넣은 뒤 자식들을 넣는다. 이후 같은 노드가 다시 꺼내질 때는 두 자식이 이미 계산된 상태다.

const stack = [[root, false]];

while (stack.length > 0) {
  const [node, visited] = stack.pop();
  if (node === null) continue;

  if (!visited) {
    stack.push([node, true]);
    stack.push([node.right, false]);
    stack.push([node.left, false]);
    continue;
  }

  // 이 시점에는 양쪽 자식의 gain이 계산되어 있다.
}

3-2. 정답 후보와 부모에게 줄 값을 분리하기

자식의 기여값은 Map에 저장한다. 음수는 버리고, 양쪽을 합친 값으로 정답을 갱신한 다음 한쪽만 선택한 값을 현재 노드의 기여값으로 기록한다.

const leftGain = Math.max(0, gain.get(node.left) ?? 0);
const rightGain = Math.max(0, gain.get(node.right) ?? 0);

answer = Math.max(answer, node.val + leftGain + rightGain);
gain.set(node, node.val + Math.max(leftGain, rightGain));

3-3. 경계 조건 확인하기

• 노드가 하나뿐이면 그 노드 값이 정답이다.
• 모든 값이 음수여도 경로는 비어 있을 수 없으므로 가장 큰 노드 값이 남는다.
• 자식의 기여값이 음수라면 0으로 취급해 경로에서 제외한다.
• 루트를 지나지 않는 최적 경로도 각 노드에서 전역 최댓값을 갱신하므로 놓치지 않는다.
• 한쪽으로 30,000개가 이어진 트리도 명시적 스택을 사용하므로 재귀 깊이에 의존하지 않는다.

4. 최종 구현

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val === undefined ? 0 : val);
 *     this.left = (left === undefined ? null : left);
 *     this.right = (right === undefined ? null : right);
 * }
 */

/**
 * @param {TreeNode} root
 * @return {number}
 */
var maxPathSum = function (root) {
  let answer = -Infinity;
  const gain = new Map();
  const stack = [[root, false]];

  while (stack.length > 0) {
    const [node, visited] = stack.pop();
    if (node === null) continue;

    if (!visited) {
      stack.push([node, true]);
      stack.push([node.right, false]);
      stack.push([node.left, false]);
      continue;
    }

    const leftGain = Math.max(0, gain.get(node.left) ?? 0);
    const rightGain = Math.max(0, gain.get(node.right) ?? 0);

    answer = Math.max(answer, node.val + leftGain + rightGain);
    gain.set(node, node.val + Math.max(leftGain, rightGain));
  }

  return answer;
};

5. 시간복잡도

노드 수를 n, 트리의 높이를 h라고 하자.

• 시간복잡도: O(n) — 각 노드를 처음 방문할 때와 처리할 때 한 번씩 확인한다.
• 공간복잡도: O(n) — 모든 노드의 기여값을 Map에 저장하고, 명시적 스택도 최악의 경우 O(n)을 사용한다.

결과적으로 부모에게 이어 줄 수 있는 한 방향의 경로와 현재 노드에서 완성되는 양방향 경로를 분리하면, 모든 경로를 직접 만들지 않고도 최대 경로 합을 구할 수 있다.

참고자료

• LeetCode: https://leetcode.com/problems/binary-tree-maximum-path-sum/
• Wikipedia — Tree traversal: https://en.wikipedia.org/wiki/Tree_traversal