ZigZag traversal

This is one of the question from algoexpert and this is my solution involving math

def zigzagTraversal(arr):
out = []
test = 1
out.append(arr[0][0])
while test <= len(arr)+len(arr[0]):
    for j in range(len(arr[0])):
        for i in range(len(arr)):
            if test % 2 == 1:
                if i + j == test:
                    # print(arr[i][j])
                    out.append(arr[i][j])
            else:
                if i + j == test:
                    out.append(arr[j][i])
    test += 1
return out
print(zigzagTraversal([[1, 3, 4, 10],
                   [2, 5, 9, 11],
                   [6, 8, 12, 15],
                   [7, 13, 14, 16]]))

Let N=len(arr), M=len(arr[0]). This code has O(NM^2+N^2M) complexity (or O(N^3) for square input matrix). Here's why: outer loop (while) is fully equivalent to for loop (because you increment test at the end and never change it elsewhere) and will be executed N+M times. Then you enter the second loop which is executed exactly N times. Then inner loop - M times. Inner loop will be repeated for each step of middle loop, the same with outer. So we have to multiply all this counts together. All inner operations are not atomic, of course, but for small data will be O(1) on average. If your input is large, I suggest pre-allocating out (creating it as out = [None for _ in range(N*M)]) to avoid costs of its growth and keeping current first free index as additional variable.