English
498. Diagonal Traverse 
Problem Statement
Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.
Example 1:
Input: mat = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,2,4,7,5,3,6,8,9]
Example 2:
Input: mat = [[1,2],[3,4]]
Output: [1,2,3,4]
Constraints:
m == mat.lengthn == mat[i].length1 <= m, n <= 1041 <= m * n <= 104-105 <= mat[i][j] <= 105
Solution:
go
package main
type Direction int
const (
Up Direction = iota
Down
)
func findDiagonalOrder(mat [][]int) []int {
direction := Up
row, col := 0, 0
row_len, col_len := len(mat), len(mat[0])
res := make([]int, row_len*col_len)
for i := 0; i < row_len*col_len; i++ {
res[i] = mat[row][col]
if direction == Up {
if row == 0 || col == col_len-1 {
direction = Down
if col == col_len-1 {
row++
} else {
col++
}
} else {
row--
col++
}
} else {
if col == 0 || row == row_len-1 {
direction = Up
if row == row_len-1 {
col++
} else {
row++
}
} else {
row++
col--
}
}
}
return res
}
rs
enum Direction {
Up,
Down,
}
impl Solution {
pub fn find_diagonal_order(mat: Vec<Vec<i32>>) -> Vec<i32> {
let mut result = vec![];
let mut direction = Direction::Up;
let (mut row, mut col) = (0, 0);
let (row_len, col_len) = (mat.len(), mat[0].len());
while row < row_len && col < col_len {
// push the current element to the result vector
result.push(mat[row][col]);
match direction {
Direction::Up => {
if row == 0 || col == col_len - 1 {
// if we are at the top row or the rightmost column, change direction to "Down"
direction = Direction::Down;
if col == col_len - 1 {
// if at the rightmost column, move to the next row
row += 1;
} else {
// otherwise, move to the next column
col += 1;
}
} else {
// move diagonally upward
row -= 1;
col += 1;
}
}
Direction::Down => {
if col == 0 || row == row_len - 1 {
// if we are at the leftmost column or the bottom row, change direction to "Up"
direction = Direction::Up;
if row == row_len - 1 {
// if at the bottom row, move to the next column
col += 1;
} else {
// otherwise, move to the next row
row += 1;
}
} else {
// move diagonally downward
row += 1;
col -= 1;
}
}
}
}
result
}
}
...