-
Notifications
You must be signed in to change notification settings - Fork 1
/
flat.go
96 lines (88 loc) · 1.68 KB
/
flat.go
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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
package flat
// Node represents a multi level doubly linked list node
type Node struct {
Val int
Prev *Node
Next *Node
Child *Node
}
// func constructFromSlice([]int) *Node {
// return nil
// }
func flatten(head *Node) *Node {
return useDFS(head)
}
// useDFS time complexity O(N), space complexity O(N)
func useDFS(head *Node) *Node {
if head == nil {
return head
}
dummy := &Node{Next: head}
dfs(dummy, head)
dummy.Next.Prev = nil
return dummy.Next
}
// dfs preOrder traversal
func dfs(prev, cur *Node) *Node {
if cur == nil {
return prev
}
// link prev and cur
prev.Next = cur
cur.Prev = prev
next := cur.Next
// left sub tree
newHead := dfs(cur, cur.Child)
cur.Child = nil
// right sub tree
return dfs(newHead, next)
}
// useStack time complexity O(N) space complexity O(N)
func useStack(head *Node) *Node {
if head == nil {
return head
}
dummy := &Node{Next: head}
stack := make([]*Node, 0, 1000)
prev, cur := dummy, dummy
stack = append(stack, head)
n := len(stack)
for n > 0 {
cur = stack[n-1]
stack = stack[:n-1]
prev.Next = cur
cur.Prev = prev
if cur.Next != nil {
// right sub tree is not nil
stack = append(stack, cur.Next)
}
if cur.Child != nil {
// left sub tree
stack = append(stack, cur.Child)
cur.Child = nil
}
prev = cur
n = len(stack)
}
return dummy.Next
}
// iterator time complexity O(N), space complexity O(1)
func iterator(head *Node) *Node {
for h := head; h != nil; h = h.Next {
if h.Child != nil {
next := h.Next
h.Next = h.Child
h.Next.Prev = h
h.Child = nil
p := h.Next
for p.Next != nil {
p = p.Next
}
p.Next = next
if next != nil {
next.Prev = p
}
}
}
return head
}