116. Populating Next Right Pointers in Each Node - Medium

前往題目

想法

• 用Breadth-First Search

思路

1. Breadth-First Search
2. 搜尋的時候建立next指針，因為每次循環都剛好是整層，疊代整層的時候順便把指針也設定好了

Code

一陣子沒寫BFS了，但還記得!

class Solution {
    public Node connect(Node root) {
        if (root == null) return null;
        
        Queue<Node> q = new LinkedList();

        q.offer(root);

        while (!q.isEmpty()) {
            int size = q.size();

            if (size == 1) {
                Node cur = q.poll();
                if (cur.left != null) q.offer(cur.left);
                if (cur.right != null) q.offer(cur.right);
                continue;
            }

            Node prevNode = q.poll();
            if (prevNode.left != null) q.offer(prevNode.left);
            if (prevNode.right != null) q.offer(prevNode.right);

            for (int i = 1; i < size; ++i) {            
                Node curNode = q.poll();
                prevNode.next = curNode;
                prevNode = curNode;

                if (curNode.left != null) q.offer(curNode.left);
                if (curNode.right != null) q.offer(curNode.right);
            }
        }

        return root;
    }
}

另外NeetCode給的解法更加的簡潔

class Solution {
    public Node connect(Node root) {
        Node current = root;
        Node next = root == null ? null : root.left;

        while (current != null && next != null) {
            // Connect the children of the current node
            current.left.next = current.right;

            // When the node besides the current node exists
            if (current.next != null) {
                current.right.next = current.next.left;
            }

            // Assign current node to its neighbor
            current = current.next;

            // When the neighbor doesn't exist
            if (current == null) {
                current = next;
                next = current.left;
            }
        }

        return root;
    }
}