forked from rui-yan/LeetCode-1
-
Notifications
You must be signed in to change notification settings - Fork 0
/
populating-next-right-pointers-in-each-node-ii.py
76 lines (67 loc) · 1.9 KB
/
populating-next-right-pointers-in-each-node-ii.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from __future__ import print_function
# Time: O(n)
# Space: O(1)
#
# 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
#
# Definition for a binary tree node
class TreeNode:
def __init__(self, x):
self.val = x
self.left = None
self.right = None
self.next = None
def __repr__(self):
if self is None:
return "Nil"
else:
return "{} -> {}".format(self.val, repr(self.next))
class Solution:
# @param root, a tree node
# @return nothing
def connect(self, root):
head = root
while head:
prev, cur, next_head = None, head, None
while cur:
if next_head is None:
if cur.left:
next_head = cur.left
elif cur.right:
next_head = cur.right
if cur.left:
if prev:
prev.next = cur.left
prev = cur.left
if cur.right:
if prev:
prev.next = cur.right
prev = cur.right
cur = cur.next
head = next_head
if __name__ == "__main__":
root, root.left, root.right = TreeNode(1), TreeNode(2), TreeNode(3)
root.left.left, root.left.right, root.right.right = TreeNode(4), TreeNode(5), TreeNode(7)
Solution().connect(root)
print(root)
print(root.left)
print(root.left.left)