LeetCode：5394. 对角线遍历 II

https://leetcode-cn.com/problems/diagonal-traverse-ii

给你一个列表 nums ，里面每一个元素都是一个整数列表。请你依照下面各图的规则，按顺序返回 nums 中对角线上的整数。

示例 1：

输入：nums = [[1,2,3],[4,5,6],[7,8,9]]
输出：[1,4,2,7,5,3,8,6,9]

示例 2：

输入：nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
输出：[1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

示例 3：

输入：nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
输出：[1,4,2,5,3,8,6,9,7,10,11]

示例 4：

输入：nums = [[1,2,3,4,5,6]]
输出：[1,2,3,4,5,6]
 

提示：

1 <= nums.length <= 10^5
1 <= nums[i].length <= 10^5
1 <= nums[i][j] <= 10^9
nums 中最多有 10^5 个数字。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/diagonal-traverse-ii
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

 

思路：注意到箭头线上的位置和(i+j)均相同，以位置和为字典key保存元素，同时保存列坐标j，按列坐标排序，输出排序后的字典即可

class Solution(object):
    def findDiagonalOrder(self, nums):
        """
        :type nums: List[List[int]]
        :rtype: List[int]
        """
        # 注意到箭头线上的位置和(i+j)均相同，以位置和为字典key保存元素，同时保存列坐标j，按列坐标排序，输出排序后的字典即可
        nums_diagonal_dict = {}
        for i, num in enumerate(nums):
            for j, num_value in enumerate(num):
                if i+j in nums_diagonal_dict:
                    nums_diagonal_dict[i+j].append((j, num_value))
                else:
                    nums_diagonal_dict[i+j] = []
                    nums_diagonal_dict[i+j].append((j, num_value))
        for key in nums_diagonal_dict:
            nums_diagonal_dict[key].sort()

        nums_diagonal = []
        for key in nums_diagonal_dict:
            for value in nums_diagonal_dict[key]:
                nums_diagonal.append(value[1])

        return nums_diagonal