Create sibling relationships within a Binary Tree in O(1) memory

generate_siblings.py

#!/usr/bin/env python
"""
Find the right sibling in a Binary Tree that's not fully populated

e.g.

        1
       / \
      2   6
     / \   \
    3   4   7
       /     \
      5       8

Siblings:

1 => None
2 => 6
3 => 4
4 => 7
5 => 8
6 => None
7 => None
8 => None
"""


class Node(object):
    """The Node Class"""
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
        self.sibling = None

    def __repr__(self):
        return 'Node({})'.format(self.val)


def construct_bt():
    """Construct the Binary Tree"""
    _5 = Node(5)
    _8 = Node(8)
    _4 = Node(4, left=_5)
    _7 = Node(7, right=_8)
    _3 = Node(3)
    _2 = Node(2, left=_3, right=_4)
    _6 = Node(6, right=_7)
    _1 = Node(1, left=_2, right=_6)
    return _1


def create_sibling_relationships(root):
    root.sibling = None
    upper_left = root
    while upper_left:
        new_upper_left = None
        node = upper_left
        next_ = None
        while node:
            if node.left:
                if not new_upper_left:
                    new_upper_left = node.left
                    next_ = new_upper_left
                else:
                    next_.sibling = node.left
                    next_ = node.left
            if node.right:
                if not new_upper_left:
                    new_upper_left = node.right
                    next_ = new_upper_left
                else:
                    next_.sibling = node.right
                    next_ = node.right
            node = node.sibling
        upper_left = new_upper_left


def main():
    root = construct_bt()
    create_sibling_relationships(root)
    print root.left.right.left.sibling

if __name__ == '__main__':
    main()