Add: Add One Row to Tree

Author: openset

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
 }

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
+}

problems/add-one-row-to-tree/add_one_row_to_tree_test.go

@@ -1 +1,40 @@
 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},
+		},
+	}
+	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)
+		}
+	}
+}