diff --git a/internal/kit/tree_node.go b/internal/kit/tree_node.go index 2f94cd22d..d746c1687 100644 --- a/internal/kit/tree_node.go +++ b/internal/kit/tree_node.go @@ -48,8 +48,6 @@ func TreeNode2SliceInt(t *TreeNode) (s []int) { s = append(s, queue[0].Val) if queue[0].Left != nil || queue[0].Right != nil { queue = append(queue, queue[0].Left) - } - if queue[0].Right != nil { queue = append(queue, queue[0].Right) } } else { @@ -58,5 +56,8 @@ func TreeNode2SliceInt(t *TreeNode) (s []int) { queue = queue[1:] } } + for i := len(s) - 1; i >= 0 && s[i] == NULL; i-- { + s = s[:i] + } return } diff --git a/problems/add-one-row-to-tree/add_one_row_to_tree.go b/problems/add-one-row-to-tree/add_one_row_to_tree.go index 1404bed3e..d74d04c0a 100644 --- a/problems/add-one-row-to-tree/add_one_row_to_tree.go +++ b/problems/add-one-row-to-tree/add_one_row_to_tree.go @@ -1 +1,31 @@ package add_one_row_to_tree + +import "github.com/openset/leetcode/internal/kit" + +type TreeNode = kit.TreeNode + +/** + * Definition for a binary tree node. + * type TreeNode struct { + * Val int + * Left *TreeNode + * Right *TreeNode + * } + */ +func addOneRow(root *TreeNode, v int, d int) *TreeNode { + switch d { + case 1: + return &TreeNode{Val: v, Left: root} + case 2: + root.Left = &TreeNode{Val: v, Left: root.Left} + root.Right = &TreeNode{Val: v, Right: root.Right} + default: + if root.Left != nil { + root.Left = addOneRow(root.Left, v, d-1) + } + if root.Right != nil { + root.Right = addOneRow(root.Right, v, d-1) + } + } + return root +} diff --git a/problems/add-one-row-to-tree/add_one_row_to_tree_test.go b/problems/add-one-row-to-tree/add_one_row_to_tree_test.go index 1404bed3e..62ef069e5 100644 --- a/problems/add-one-row-to-tree/add_one_row_to_tree_test.go +++ b/problems/add-one-row-to-tree/add_one_row_to_tree_test.go @@ -1 +1,52 @@ package add_one_row_to_tree + +import ( + "reflect" + "testing" + + "github.com/openset/leetcode/internal/kit" +) + +type caseType struct { + input []int + v int + d int + expected []int +} + +func TestAddOneRow(t *testing.T) { + tests := [...]caseType{ + { + input: []int{4, 2, 6, 3, 1, 5}, + v: 1, + d: 2, + expected: []int{4, 1, 1, 2, kit.NULL, kit.NULL, 6, 3, 1, 5}, + }, + { + input: []int{4, 2, kit.NULL, 3, 1}, + v: 1, + d: 3, + expected: []int{4, 2, kit.NULL, 1, 1, 3, kit.NULL, kit.NULL, 1}, + }, + { + input: []int{3, 1}, + v: 1, + d: 1, + expected: []int{1, 3, kit.NULL, 1}, + }, + { + input: []int{3, kit.NULL, 2}, + v: 1, + d: 3, + expected: []int{3, kit.NULL, 2, 1, 1}, + }, + } + for _, tc := range tests { + root := kit.SliceInt2TreeNode(tc.input) + root = addOneRow(root, tc.v, tc.d) + output := kit.TreeNode2SliceInt(root) + if !reflect.DeepEqual(output, tc.expected) { + t.Fatalf("input: %v, output: %v, expected: %v", tc.input, output, tc.expected) + } + } +} diff --git a/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree.go b/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree.go index fb4d9e461..8f9366adf 100644 --- a/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree.go +++ b/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree.go @@ -1,6 +1,8 @@ package convert_sorted_array_to_binary_search_tree -import . "github.com/openset/leetcode/internal/kit" +import "github.com/openset/leetcode/internal/kit" + +type TreeNode = kit.TreeNode /** * Definition for a binary tree node. diff --git a/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree_test.go b/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree_test.go index 855ae9855..2d18a69f7 100644 --- a/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree_test.go +++ b/problems/convert-sorted-array-to-binary-search-tree/convert_sorted_array_to_binary_search_tree_test.go @@ -4,7 +4,7 @@ import ( "reflect" "testing" - . "github.com/openset/leetcode/internal/kit" + "github.com/openset/leetcode/internal/kit" ) type caseType struct { @@ -16,11 +16,11 @@ func TestSortedArrayToBST(t *testing.T) { tests := [...]caseType{ { input: []int{-10, -3, 0, 5, 9}, - expected: []int{0, -3, 9, -10, 5}, + expected: []int{0, -3, 9, -10, kit.NULL, 5}, }, } for _, tc := range tests { - output := TreeNode2SliceInt(sortedArrayToBST(tc.input)) + output := kit.TreeNode2SliceInt(sortedArrayToBST(tc.input)) if !reflect.DeepEqual(output, tc.expected) { t.Fatalf("input: %v, output: %v, expected: %v", tc.input, output, tc.expected) }