leetcode_rust/
diagonal_traverse.rs

1#![allow(dead_code)]
2
3/// When meet following conditions:
4/// 1. up and j == n-1: i++
5/// 2. up and i == 0: j++
6/// 3. down and j == m-1: j++
7/// 4. down and j == 0: i++
8/// and let up = !up
9/// otherwise go up and right or go down and left.
10pub fn find_diagonal_order(matrix: Vec<Vec<i32>>) -> Vec<i32> {
11    if matrix.is_empty() {
12        return vec![];
13    }
14    let mut i = 0;
15    let mut j = 0;
16    let mut res = vec![];
17    let m = matrix.len();
18    let n = matrix[0].len();
19    let mut up = true;
20    let level = m + n - 2;
21    while i + j < level {
22        res.push(matrix[i][j]);
23
24        if up {
25            if j == n - 1 {
26                i += 1;
27                up = !up;
28            } else if i == 0 {
29                j += 1;
30                up = !up;
31            } else {
32                i -= 1;
33                j += 1;
34            }
35        } else {
36            if i == m - 1 {
37                j += 1;
38                up = !up;
39            } else if j == 0 {
40                i += 1;
41                up = !up;
42            } else {
43                i += 1;
44                j -= 1;
45            }
46        }
47    }
48    res.push(matrix[i][j]);
49
50    res
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn test1() {
59        assert_eq!(
60            find_diagonal_order(vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]]),
61            vec![1, 2, 4, 7, 5, 3, 6, 8, 9]
62        );
63
64        assert_eq!(
65            find_diagonal_order(vec![vec![1, 2], vec![3, 4], vec![5, 6]]),
66            vec![1, 2, 3, 5, 4, 6]
67        );
68    }
69}
leetcode_rust/diagonal_traverse.rs

leetcode_rust/
diagonal_traverse.rs
