Populating Next Right Pointers in Each Node II

题目描述

Follow up for problem "Populating Next Right Pointers in Each Node".

What if the given tree could be any binary tree? Would your previous solution still work?

Note:

You may only use constant extra space. For example, Given the following binary tree,

         1
       /  \
      2    3
     / \    \
    4   5    7

After calling your function, the tree should look like:

         1 -> NULL
       /  \
      2 -> 3 -> NULL
     / \    \
    4-> 5 -> 7 -> NULL

解题方法

不再是完全二叉树,但是bfs level order的code依然work

Solution

from collections import deque
# Definition for binary tree with next pointer.
# class TreeLinkNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
#         self.next = None

class Solution(object):
    def connect(self, root):
        """
        :type root: TreeLinkNode
        :rtype: nothing
        """
        if not root:
            return
        queue = deque()
        result = []
        queue.append(root)
        while queue:
            size = len(queue)
            for i in range(size):
                cur = queue.popleft()
                if i == size - 1:
                    cur.next = None
                else:
                    cur.next = queue[0]
                if cur.left:
                    queue.append(cur.left)
                if cur.right:
                    queue.append(cur.right)

O(1)

from collections import deque
# Definition for binary tree with next pointer.
# class TreeLinkNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
#         self.next = None

class Solution(object):
    def connect(self, root):
        """
        :type root: TreeLinkNode
        :rtype: nothing
        """
        if not root:
            return
        # prev level's head
        lastHead = root
        # prev level's pointer
        lastCurrent = None
        # current level's head
        curHead = None
        # current level's pointer
        cur = None
        while lastHead:
            lastCurrent = lastHead # start another level loop
            while lastCurrent:
                if lastCurrent.left:
                    if curHead is None:
                        # start from the beginning
                        curHead = lastCurrent.left
                        cur = lastCurrent.left
                    else:
                        cur.next = lastCurrent.left
                        cur = cur.next
                if lastCurrent.right:
                    if curHead is None:
                        # start from the beginning
                        curHead = lastCurrent.right
                        cur = lastCurrent.right
                    else:
                        cur.next = lastCurrent.right
                        cur = cur.next
                lastCurrent = lastCurrent.next
            lastHead = curHead
            curHead = None
public void connect(TreeLinkNode root) {
    if(root == null) 
        return;

    TreeLinkNode lastHead = root;//prevous level's head 
    TreeLinkNode lastCurrent = null;//previous level's pointer
    TreeLinkNode currentHead = null;//currnet level's head 
    TreeLinkNode current = null;//current level's pointer

    while(lastHead!=null){
        lastCurrent = lastHead;

        while(lastCurrent!=null){
            //left child is not null
            if(lastCurrent.left!=null)    {
                if(currentHead == null){
                    currentHead = lastCurrent.left;
                    current = lastCurrent.left;
                }else{
                    current.next = lastCurrent.left;
                    current = current.next;
                }
            }

            //right child is not null
            if(lastCurrent.right!=null){
                if(currentHead == null){
                    currentHead = lastCurrent.right;
                    current = lastCurrent.right;
                }else{
                    current.next = lastCurrent.right;
                    current = current.next;
                }
            }

            lastCurrent = lastCurrent.next;
        }

        //update last head
        lastHead = currentHead;
        currentHead = null;
    }
}

Reference