[leetcode #116] Populating Next Right Pointers in Each Node

Problem

You are given a perfect binary tree where all leaves are on the same level, and every parent has two children. The binary tree has the following definition:

struct Node {
  int val;
  Node *left;
  Node *right;
  Node *next;
}

Populate each next pointer to point to its next right node. If there is no next right node, the next pointer should be set to NULL.

Initially, all next pointers are set to NULL.

Example 1:

Input: root = [1,2,3,4,5,6,7]
Output: [1,#,2,3,#,4,5,6,7,#]
Explanation: Given the above perfect binary tree (Figure A), your function should populate each next pointer to point to its next right node, just like in Figure B. The serialized output is in level order as connected by the next pointers, with '#' signifying the end of each level.

Example 2:

Input: root = []
Output: []

Constraints:

・ The number of nodes in the tree is in the range [0, 212 - 1].
・ -1000 <= Node.val <= 1000

Follow-up:

・ You may only use constant extra space.
・ The recursive approach is fine. You may assume implicit stack space does not count as extra space for this problem.

Idea

tree 문제는 항상 재미있다. 평소엔 잘 활용하지 않던 recursive function을 구현하는 것은 언제나 흥미로운 주제이기 때문이다.

이번 문제는 같은 level에 있는 node를 왼쪽부터 오른쪽으로 연결하는 문제다.

map을 이용하면 쉽게 풀 수 있다. Tree를 탐색하면서 같은 level에 있는 리스트에 각 node를 더한다. Tree 탐색이 끝난 뒤에 list를 순차적으로 탐색하면서 같은 level에 있는 노드의 next pointer를 하나씩 바꿔주면 된다.

위 방법은 space complexity가 O(n)이다. 그런데 space complexity를 O(1)로 해서 풀 수 있는 방법이 있다.

Tree를 탐색하면서 next pointer가 존재할 경우 현재 노드의 오른쪽 자식 노드의 next pointer를 next 노드의 왼쪽 자식 노드를 가리키게 만드는 것이다.

위와 같이 탐색을 반복하면 같은 레벨에 있는 노드들이 next pointer로 연결되게 된다.

Solution

Map을 이용해 푸는 방법은 아래와 같다.

/*
// Definition for a Node.
class Node {
    public int val;
    public Node left;
    public Node right;
    public Node next;

    public Node() {}

    public Node(int _val) {
        val = _val;
    }

    public Node(int _val, Node _left, Node _right, Node _next) {
        val = _val;
        left = _left;
        right = _right;
        next = _next;
    }
};
*/

class Solution {
    Map<Integer, List<Node>> map = new HashMap<>();

    public Node connect(Node root) {
        traverse(root, 0);

        for (List<Node> list : map.values()) {
            Node node = list.get(0);
            for (int i=1; i < list.size(); i++) {
                node.next = list.get(i);
                node = list.get(i);
            }
        }

        return root;
    }

    private void traverse(Node node, int level) {
        if (node == null) {
            return;
        }

        List<Node> nodeList = map.getOrDefault(level, new ArrayList<>());
        nodeList.add(node);
        map.put(level, nodeList);

        traverse(node.left, level+1);
        traverse(node.right, level+1);
    } 
}

Map 없이 space complexity를 O(1)로 푸는 법은 다음과 같다.

class Solution {    
    public Node connect(Node root) {
        if (root == null) {
            return null;
        }

        if (root.left != null) {
            root.left.next = root.right;
            if (root.next != null) {
                root.right.next = root.next.left;
            }
        }

        connect(root.left);
        connect(root.right);

        return root;
    }
}

Reference

https://leetcode.com/problems/populating-next-right-pointers-in-each-node/