LeetCode #6: zigzag 处理边界条件,测试通过

go/cmd/demo.go

@@ -23,10 +23,9 @@ var demoCmd = &cobra.Command{
 // utf8.RuneCountInString("Hi")
 
 func demoMain() {
-	// demo.Convert("A", 1)
+	fmt.Println(demo.Convert("AB", 1))
+	fmt.Println(demo.Convert("A", 1))
 	demo.Convert("ABC", 2)
 	r := demo.Convert("PAYPALISHIRING", 3)
 	fmt.Println(r == "PAHNAPLSIIGYIR", r)
-	fmt.Println(r)
-	fmt.Println("PAHNAPLSIIGYIR")
 }

go/demo/leetcode.go

@@ -272,8 +272,24 @@ func Convert(s string, numRows int) string {
 	ti := 0
 	for row := 0; row < numRows; row++ {
 		o1, o2 := start-(row*2), row*2
+		// fmt.Printf("o1 %d, o2: %d of row %d\n", o1, o2, row)
 
-		if row == 0 || row == numRows-1 {
+		if row == 0 {
+			o := Max(o1, o2)
+			if o == 0 {
+				o = 1
+			}
+			si := row
+			for i := 0; ; i++ {
+				if si >= l {
+					break
+				}
+				t[ti] = r[si]
+				ti++
+
+				si += o
+			}
+		} else if row == numRows-1 {
 			o := Max(o1, o2)
 			si := row
 			for i := 0; ; i++ {

go/demo/leetcode_test.go

@@ -84,3 +84,14 @@ func TestIsPalindrome(t *testing.T) {
 	fn("abcba", true)
 	fn("abcdcba", true)
 }
+
+func TestConvert(t *testing.T) {
+	fn := func(s string, nr int, r string) {
+		if Convert(s, nr) != r {
+			t.Fail()
+		}
+	}
+	fn("A", 1, "A")
+	fn("PAYPALISHIRING", 3, "PAHNAPLSIIGYIR")
+	fn("AB", 1, "AB")
+}