1424. Diagonal Traverse II

Given a list of lists of integers, 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]

Example 3:

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

Example 4:

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

 

Constraints:

• 1 <= nums.length <= 10^5
• 1 <= nums[i].length <= 10^5
• 1 <= nums[i][j] <= 10^9
• There at most 10^5 elements in nums.

思路：同一个斜方向的row+col等于同一个数

class Solution(object):
    def findDiagonalOrder(self, nums):
        """
        :type nums: List[List[int]]
        :rtype: List[int]
        """
        from collections import defaultdict
        n = len(nums)
        m = max(len(s) for s in nums)
        ma = n-1+m-1

        d = defaultdict(list)
        for i,a in enumerate(nums):
            for j,b in enumerate(a):
                d[i+j].append(b)

        res = []
        for s in range(ma+1):
            res += d[s][::-1]
        return res


s=Solution()
print(s.findDiagonalOrder([[1,2,3],[4,5,6],[7,8,9]]))
print(s.findDiagonalOrder([[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]))