struct Solution;

impl Solution {
    fn find_diagonal_order(nums: Vec<Vec<i32>>) -> Vec<i32> {
        let n = nums.len();
        let mut rows: Vec<Vec<i32>> = vec![];
        for i in 0..n {
            for (j, &v) in nums[i].iter().enumerate() {
                let k = i + j;
                if rows.len() == k {
                    rows.push(vec![]);
                }
                rows[k].push(v);
            }
        }
        rows.into_iter().flat_map(|q| q.into_iter().rev()).collect()
    }
}

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