116. Populating Next Right Pointers in Each Node

Description

Given a binary tree, the goal is to populate each node's next pointer to point to its next right node. If there is no next right node, the next pointer should be set to null. The tree can have varying levels of depth, and some nodes might have one or no children.

Intuition

To connect nodes at the same level, a breadth-first traversal of the tree is ideal because it allows us to process each level of the tree separately. By processing nodes level by level, we can easily link them together using the next pointer.

Approach

1. Edge Case Handling: - If the root is null or the tree has no children, return the root immediately.
2. Breadth-First Search (BFS): - Use a queue to perform a BFS on the tree. The queue will store pairs of nodes and their corresponding levels. - Initialize a map (prevNodeByLevel) to track the last node visited at each level.
3. Processing Nodes: - For each node processed, check if there was a previous node at the same level using the map. - If a previous node exists, set its next pointer to the current node. - Update the map to set the current node as the last node visited at this level.
4. Child Nodes: - If the current node has left or right children, add them to the queue with their corresponding level incremented.
5. Return the Root: - After processing all nodes, return the modified tree's root.

Complexity

Time Complexity:

(O(N)), where (N) is the number of nodes in the tree. Each node is visited once.

Space Complexity:

(O(N)), where (N) is the number of nodes. This accounts for the space used by the queue during BFS and the map used to store the last node at each level.

Code

C++

/*
class Node {
public:
    int val;
    Node* left;
    Node* right;
    Node* next;

    Node() : val(0), left(NULL), right(NULL), next(NULL) {}
    Node(int _val) : val(_val), left(NULL), right(NULL), next(NULL) {}
    Node(int _val, Node* _left, Node* _right, Node* _next)
        : val(_val), left(_left), right(_right), next(_next) {}
};
*/
class Solution {
public:
    Node* connect(Node* root) {
        if (!root) return root;

        queue> q;
        q.push({root, 0});
        unordered_map prevNodeByLevel;

        while (!q.empty()) {
            auto [node, level] = q.front();
            q.pop();

            if (prevNodeByLevel.count(level)) {
                prevNodeByLevel[level]->next = node;
            }

            if (node->left) {
                q.push({node->left, level + 1});
            }

            if (node->right) {
                q.push({node->right, level + 1});
            }

            prevNodeByLevel[level] = node;
        }

        return root;
    }
};

Python

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        if not root:
            return root

        queue = [(root, 0)]
        prev_node_by_level = {}

        while queue:
            node, level = queue.pop(0)

            if level in prev_node_by_level:
                prev_node_by_level[level].next = node

            if node.left:
                queue.append((node.left, level + 1))

            if node.right:
                queue.append((node.right, level + 1))

            prev_node_by_level[level] = node

        return root

Java

/*
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 {
    public Node connect(Node root) {
        if (root == null) return root;

        Queue> queue = new LinkedList<>();
        Map prevNodeByLevel = new HashMap<>();
        queue.add(new Pair<>(root, 0));

        while (!queue.isEmpty()) {
            Pair pair = queue.poll();
            Node node = pair.getKey();
            int level = pair.getValue();

            if (prevNodeByLevel.containsKey(level)) {
                prevNodeByLevel.get(level).next = node;
            }

            if (node.left != null) {
                queue.add(new Pair<>(node.left, level + 1));
            }

            if (node.right != null) {
                queue.add(new Pair<>(node.right, level + 1));
            }

            prevNodeByLevel.put(level, node);
        }

        return root;
    }
}

JavaScript

var connect = function (root) {
    
    if (!root || (!root.left && !root.right)) {
        return root;
    }

    let queue = [[root, 0]];
    let prevNodeByLevel = new Map();

    while (queue.length) {
        let [node, level] = queue.shift();

        if (prevNodeByLevel.has(level)) {
            let prevNode = prevNodeByLevel.get(level);
            prevNode.next = node;
        }

        if (node.left) {
            queue.push([node.left, level + 1]);
        }

        if (node.right) {
            queue.push([node.right, level + 1]);
        }

        prevNodeByLevel.set(level, node);
    }

    return root;
};