498. Diagonal Traverse

Author: vfedorovatlas

diagonal-traverse/Readme.md

@@ -0,0 +1,4 @@
+# 498. Diagonal Traverse
+
+Source: https://leetcode.com/problems/diagonal-traverse/description/
+

diagonal-traverse/main.go

@@ -0,0 +1,41 @@
+package leecode
+
+func findDiagonalOrder(matrix [][]int) []int {
+	m := len(matrix)
+	if m == 0 {
+		return []int{}
+	}
+	n := len(matrix[0])
+	if m == 1 && n == 1 {
+		return matrix[0]
+	}
+	ret := []int{}
+	var row, col int
+	for i := 0; i < m*n; i++ {
+		ret = append(ret, matrix[row][col])
+		if (row+col)%2 == 0 { // moving up
+			if col == n-1 {
+				row++
+				continue
+			}
+			if row == 0 {
+				col++
+				continue
+			}
+			row--
+			col++
+		} else { // moving down
+			if row == m-1 {
+				col++
+				continue
+			}
+			if col == 0 {
+				row++
+				continue
+			}
+			row++
+			col--
+		}
+	}
+	return ret
+}

diagonal-traverse/main_test.go

@@ -0,0 +1,41 @@
+package leecode
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+func TestFindDiagonalOrder(t *testing.T) {
+	testCases := []struct {
+		Input  [][]int
+		Output []int
+	}{
+		{
+			Input:  [][]int{},
+			Output: []int{},
+		},
+		{
+			Input: [][]int{
+				{1},
+			},
+			Output: []int{1},
+		},
+		{
+			Input: [][]int{
+				{1, 2, 3},
+				{4, 5, 6},
+				{7, 8, 9},
+			},
+			Output: []int{1, 2, 4, 7, 5, 3, 6, 8, 9},
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("%v:%d", tc.Input, tc.Output), func(t *testing.T) {
+			actual := findDiagonalOrder(tc.Input)
+			if !reflect.DeepEqual(actual, tc.Output) {
+				t.Errorf("expected %d, but get %d", tc.Output, actual)
+			}
+		})
+	}
+}