leetcode-rust/src/leetcode/challenge/spiral_order.rs at main · martwz/leetcode-rust · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
// 54. Spiral Matrix, Medium
// https://leetcode.com/problems/spiral-matrix/
impl Solution {
    pub fn spiral_order(matrix: Vec<Vec<i32>>) -> Vec<i32> {
        let mut ans = vec![];
        // edge cases:
        if matrix.is_empty() || matrix[0].is_empty() {
            return ans;
        } else if matrix[0].len() == 1 {
            matrix.iter().flatten().for_each(|f| {
                ans.push(*f);
            });
            return ans;
        }

        let [mut n, mut m] = [matrix.len() - 1, matrix[0].len() - 1];
        let cells = (n + 1) * (m + 1);
        let mut i = 0;

        #[derive(PartialEq)]
        enum Direction {
            RIGHT,
            DOWN,
            LEFT,
            UP,
        }

        let mut direction = Direction::RIGHT;
        let mut pos = (0, 0);
        ans.push(matrix[pos.0][pos.1]);

        while i < cells - 1 {
            if m > 0 && direction == Direction::RIGHT {
                for _j in 0..m {
                    pos = (pos.0, pos.1 + 1);
                    ans.push(matrix[pos.0][pos.1]);
                }
                if i != 0 {
                    m -= 1;
                }
                direction = Direction::DOWN;
            } else if n > 0 && direction == Direction::DOWN {
                for _j in 0..n {
                    pos = (pos.0 + 1, pos.1);
                    ans.push(matrix[pos.0][pos.1]);
                }
                n -= 1;
                direction = Direction::LEFT;
            } else if m > 0 && direction == Direction::LEFT {
                for _j in 0..m {
                    pos = (pos.0, pos.1 - 1);
                    ans.push(matrix[pos.0][pos.1]);
                }
                m -= 1;
                direction = Direction::UP;
            } else if n > 0 && direction == Direction::UP {
                for _j in 0..n {
                    pos = (pos.0 - 1, pos.1);
                    ans.push(matrix[pos.0][pos.1]);
                }
                n -= 1;
                direction = Direction::RIGHT;
            }

            i += 1;
        }

        ans
    }
}

struct Solution {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_spiral_order() {
        assert_eq!(Solution::spiral_order(vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]]), vec![1, 2, 3, 6, 9, 8, 7, 4, 5]);
    }

    #[test]
    fn test_spiral_order2() {
        assert_eq!(
            Solution::spiral_order(vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]]),
            vec![1, 2, 3, 4, 8, 12, 11, 10, 9, 5, 6, 7]
        );
    }

    #[test]
    fn test_spiral_order3() {
        assert_eq!(Solution::spiral_order(vec![vec![1]]), vec![1]);
    }

    #[test]
    fn test_spiral_order4() {
        assert_eq!(Solution::spiral_order(vec![vec![3], vec![2]]), vec![3, 2]);
    }

    #[test]
    fn test_spiral_order5() {
        assert_eq!(Solution::spiral_order(vec![vec![3, 2], vec![0, 1]]), vec![3, 2, 1, 0]);
    }

    #[test]
    fn test_spiral_order6() {
        assert_eq!(Solution::spiral_order(vec![]), vec![]);
    }

    #[test]
    fn test_spiral_order7() {
        assert_eq!(
            Solution::spiral_order(vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12], vec![13, 14, 15, 16]]),
            vec![1, 2, 3, 4, 8, 12, 16, 15, 14, 13, 9, 5, 6, 7, 11, 10]
        );
    }
}