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.

---

如上图所示，给定一个残缺矩阵，要求遍历每个对角线，输出所有元素的值。

朴素解法会TLE，特别是如果有某一行或某一列比其他行列长特别多的情况。

提示解法：观察规律发现，处于同一对角线的元素的行列下标之和相等。所以我们可以把所有元素的行列下标进行排序，如果下标和相等，则处于同一对角线，按行下标从大到小排序；如果下标和不相等，则按下标和从小到大排序。完整代码如下：

bool cmp(const pair<int, int> &p1, const pair<int, int> &p2) {
	int sum1 = p1.first + p1.second, sum2 = p2.first + p2.second;
	if (sum1 == sum2) {
		return p1.first > p2.first;
	}
	else {
		return sum1 < sum2;
	}
}
class Solution {
public:
	vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
		vector<pair<int, int>> idxs;
		for (int i = 0; i < nums.size(); ++i) {
			for (int j = 0; j < nums[i].size(); ++j) {
				idxs.push_back(make_pair(i, j));
			}
		}
		sort(idxs.begin(), idxs.end(), cmp);
		vector<int> ans;
		for (int i = 0; i < idxs.size(); ++i) {
			ans.push_back(nums[idxs[i].first][idxs[i].second]);
		}
		return ans;
	}
};

本代码提交AC，用时500MS。