Как исправить мой код для диагональной матрицы траверс (альтернативная версия)?

Question

У меня есть следующий код, который используется для альтернативной печати диагональной матрицы.Пожалуйста, посмотрите на мой код и скажите, что с ним не так?

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        # No. of rows
        m=len(matrix)
        # No. of columns
        n=len(matrix[0])
        # Counter to alternatively switch the direction. 
        up=True
        # list to store the result
        l=[]
        i,j,count=0,0,0
        # loop for traversing all the elements.
        while(count<=m*n):
            if(up):
                while(i>=0 and j<n):
                    l.append(matrix[i][j])
                    count+=1
                    i-=1
                    j+=1
                if(i<0 and j<n):
                    i=0
                if(j==n):
                    i=i+2
                    j=j-1
            else:
                while(j>=0 and i<m):
                    l.append(matrix[i][j])
                    count+=1
                    j-=1
                    i+=1
                if(j<0 and i<m):
                    j=0
                if(i==m):
                    j=j+2
                    i=i-1
            up= not up
        print(l)

Ввод:

[
   [1,2,3],
   [4,5,6],
   [7,8,9]
]

Ожидаемый ответ:

[1,2,4,7,5,3,6,8,9]

Фактический ответ:

Строка 22: Ошибка индекса: индекс списка выходит за пределы диапазона

Answer 1

Ваш код работает, просто измените ваш while (count <= m * n) на while (count < m * n).

def findDiagonalOrder(matrix):
    # No. of rows
    m = len(matrix)
    # No. of columns
    n = len(matrix[0])
    # Counter to alternatively switch the direction. 
    up = True
    # list to store the result
    l = []
    i, j, count = 0, 0, 0
    # loop for traversing all the elements.
    while (count < m * n):
        if (up):
            while (i >= 0 and j < n):
                l.append(matrix[i][j])
                count += 1
                i -= 1
                j += 1
            if (i < 0 and j < n):
                i = 0
            if (j == n):
                i = i + 2
                j = j - 1
        else:
            while (j >= 0 and i < m):
                l.append(matrix[i][j])
                count += 1
                j -= 1
                i += 1
            if (j < 0 and i < m):
                j = 0
            if (i == m):
                j = j + 2
                i = i - 1
        up = not up
    print(l)

matrix = [
   [1,2,3],
   [4,5,6],
   [7,8,9]
]

findDiagonalOrder(matrix)
#[1, 2, 4, 7, 5, 3, 6, 8, 9]