diff --git a/leetcode/others/connect.go b/leetcode/others/connect.go new file mode 100644 index 00000000..64361dca --- /dev/null +++ b/leetcode/others/connect.go @@ -0,0 +1,34 @@ +package lothers + +type Node struct { + Val int + Left *Node + Right *Node + Next *Node +} + +func connect(root *Node) *Node { + if root == nil { + return nil + } + + current := root + for current != nil { + dummy := &Node{} + tail := dummy + + for node := current; node != nil; node = node.Next { + if node.Left != nil { + tail.Next = node.Left + tail = tail.Next + } + if node.Right != nil { + tail.Next = node.Right + tail = tail.Next + } + } + current = dummy.Next + } + + return root +} diff --git a/leetcode/others/connect_test.go b/leetcode/others/connect_test.go new file mode 100644 index 00000000..f69f7a39 --- /dev/null +++ b/leetcode/others/connect_test.go @@ -0,0 +1,79 @@ +package lothers + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestConnect(t *testing.T) { + tests := []struct { + name string + root *Node + expected []interface{} + }{ + { + name: "Example 1", + root: &Node{ + Val: 1, + Left: &Node{ + Val: 2, + Left: &Node{ + Val: 4, + }, + Right: &Node{ + Val: 5, + }, + }, + Right: &Node{ + Val: 3, + Right: &Node{ + Val: 7, + }, + }, + }, + expected: []interface{}{1, "#", 2, 3, "#", 4, 5, 7, "#"}, + }, + { + name: "Example 2", + root: nil, + expected: []interface{}{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + root := connect(tt.root) + + // Serialize the tree and check the result + result := serializeTree(root) + assert.Equal(t, tt.expected, result) + }) + } +} + +func serializeTree(root *Node) []interface{} { + if root == nil { + return []interface{}{} + } + queue := []*Node{root} + var result []interface{} + for len(queue) > 0 { + levelSize := len(queue) + for i := 0; i < levelSize; i++ { + node := queue[i] + result = append(result, node.Val) + if i == levelSize-1 { + result = append(result, "#") + } + if node.Left != nil { + queue = append(queue, node.Left) + } + if node.Right != nil { + queue = append(queue, node.Right) + } + } + queue = queue[levelSize:] + } + return result +}