Diagonal Traverse II

Problem

Given a 2D integer array nums, return all elements of nums in diagonal order as shown in the below images.

 

Example 1:

Input: nums = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,4,2,7,5,3,8,6,9]

Example 2:

Input: nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
Output: [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

 

Constraints:

• 1 <= nums.length <= 105
• 1 <= nums[i].length <= 105
• 1 <= sum(nums[i].length) <= 105
• 1 <= nums[i][j] <= 105

Solution

/**
 * @param {number[][]} nums
 * @return {number[]}
 */
var findDiagonalOrder = function(nums) {
    let buckets = [];
    let n = [];

    for (let i = 0; i < nums.length; ++i) {
        for (let j = 0; j < nums[i].length; ++j) {
            let index = i + j;
            while (buckets.length <= index) {
                buckets.push([]);
            }
            buckets[index].push(nums[i][j]);
        }
    }

    for (let num of buckets) {
        for (let i = num.length - 1; i >= 0; --i) {
            n.push(num[i]);
        }
    }
    return n;
};

Notice that in a 2d array, the index of the row plus the index of the column is equal to the diagonal it is a part of. Create a 2d array buckets that stores the elements in each diagonal. Iterate through each entry in nums (row by row where i is the row index and j is the column index) and push each element to the array at buckets[i + j]. Next, flatten buckets and reverse each child array (since the values are in backwards order).