Medium? Seriously?
Well again, always smart solutions out there: https://discuss.leetcode.com/topic/77889/3-line-python-solution
class Solution(object): def findDiagonalOrder(self, matrix): r = [] n = len(matrix) if n == 0: return r m = len(matrix[0]) ni = n + m - 1 for i in range(ni): rr = [] x = min(n - 1, i) y = max(i - n + 1, 0) while x >= 0 and y < m: rr += [matrix[x][y]] x = x - 1 y = y + 1 if i % 2 == 1: rr.reverse() r = r + rr return r