1
# Runtime: 305 ms (Top 5.46%) | Memory: 15.4 MB (Top 49.03%)
2
"""
3
# Definition for a Node.
4
class Node(object):
5
    def __init__(self, val=0, left=None, right=None, next=None):
6
        self.val = val
7
        self.left = left
8
        self.right = right
9
        self.next = next
10
"""
11

12

13
class Solution(object):
14
    def findRightMost(self, root, level, requiredLevel):
15
        if not root:
16
            return root
17
        if level == requiredLevel:
18
            return root
19
        right = self.findRightMost(root.right, level + 1, requiredLevel)
20
        if right:
21
            return right
22
        return self.findRightMost(root.left, level + 1, requiredLevel)
23

24
    def findLeftMost(self, root, level, requiredLevel):
25
        if not root:
26
            return root
27
        if level == requiredLevel:
28
            return root
29
        left = self.findLeftMost(root.left, level + 1, requiredLevel)
30
        if left:
31
            return left
32
        return self.findLeftMost(root.right, level + 1, requiredLevel)
33

34
    def findRightMostFromRoot(self, rootLevelInfo, requiredLevel, currentRight):
35
        if currentRight:
36
            if currentRight.right:
37
                return currentRight.right
38
            if currentRight.left:
39
                return currentRight.left
40
        root, rootlevel = rootLevelInfo
41
        rightMost = self.findRightMost(root, rootlevel, requiredLevel)
42
        while not rightMost and root:
43
            root = root.right if root.right else root.left
44
            rootlevel += 1
45
            rightMost = self.findRightMost(root, rootlevel, requiredLevel)
46
            if rightMost:
47
                rootLevelInfo[-1] = rootlevel
48
                rootLevelInfo[0] = root
49
        return rightMost
50

51
    def findLeftMostFromRoot(self, rootLevelInfo, requiredLevel, currentLeft):
52
        if currentLeft:
53
            if currentLeft.left:
54
                return currentLeft.left
55
            if currentLeft.right:
56
                return currentLeft.right
57
        root, rootlevel = rootLevelInfo
58
        leftMost = self.findLeftMost(root, rootlevel, requiredLevel)
59
        while not leftMost and root:
60
            root = root.left if root.left else root.right
61
            rootlevel += 1
62
            leftMost = self.findLeftMost(root, rootlevel, requiredLevel)
63
            if leftMost:
64
                rootLevelInfo[-1] = rootlevel
65
                rootLevelInfo[0] = root
66

67
        return leftMost
68

69
    def stitch(self, root):
70
        if not root:
71
            return
72
        leftRootStart = [root.left, 1]
73
        rightRootStart = [root.right, 1]
74
        connectLevel = 1
75
        currentLeft = self.findLeftMostFromRoot(rightRootStart, 1, None)
76
        currentRight = self.findRightMostFromRoot(leftRootStart, 1, None)
77
        while currentLeft and currentRight:
78
            currentRight.next = currentLeft
79
            connectLevel += 1
80
            currentLeft = self.findLeftMostFromRoot(
81
                rightRootStart, connectLevel, currentLeft
82
            )
83
            currentRight = self.findRightMostFromRoot(
84
                leftRootStart, connectLevel, currentRight
85
            )
86

87
        self.stitch(root.left)
88
        self.stitch(root.right)
89

90
    def connect(self, root):
91
        """
92
        :type root: Node
93
        :rtype: Node
94
        """
95
        if not root:
96
            return root
97
        self.stitch(root)
98
        return root