Print diagonal numbers of matrix in ascending order Today I attempted an interview problem:


  Q: An n×n matrix contains numbers from 1 to n×n and is generated in the 
  fashion:
  
  The first row is of numbers starting from 1 to n left to right, Second row is from n+1 to 2n right to left, 3rd row is from 2n+1 to 3n left to right and so on.
  
  Example n = 4
  
  Matrix would look something like 

 1  2  3  4
 8  7  6  5
 9 10 11 12
16 15 14 13

  
  Diagonal elements: 1, 4, 6, 7, 10, 11, 13, 16.


Wrote this code, got an feedback as space and time complexity could be further improved.
How to improve it?

def create_matrix(n):
    mat = []
    num = 1
    reverse_flag = True
    for i in range(n):
        row = []
        reverse_flag = False if reverse_flag else True
        for j in range(n):
            row.append(num)
            num += 1                
        if reverse_flag:
            row.reverse()
        mat.append(row)
    print_matrix(mat, n)

def print_matrix(mat, n):
    reverse_flag = True
    for i in range(n):
        reverse_flag = False if reverse_flag else True
        if reverse_flag:
            if i == n-1-i:
                print mat[i][i]
            else:
                print mat[i][n-1-i], mat[i][i]
        else:
            if i == n-1-i:
                print mat[i][i]
            else:
                print mat[i][i], mat[i][n-1-i]

create_matrix(4)
It's bizarre that create_matrix() calls print_matrix().  According the Single Responsibility Principle, each function should do just one thing.  And if the function is called create_matrix(), then it should just return a matrix, and do nothing else.

If your create_matrix() function already implements the boustrophedon arrangement, then I see no reason why print_matrix() should also have a toggling reverse_flag.  That is, I consider the output


1 4
6 7
11 10
16 13



… to be wrong, as it is neither sorted in ascending order (1, 4, 6, 7, 10, 11, 13, 16) nor alternating (1, 4, 7, 6, 10, 11, 16, 13).

You don't actually have to create the matrix to print the diagonal entries.  Just use arithmetic.

def matrix_diagonals(n):
    for row_num in xrange(n):
        row_min = n * row_num + 1
        row_max = n * row_num + n
        diag1 = row_min + row_num
        diag2 = row_max - row_num
        yield max(diag1, diag2) if row_num % 2 else min(diag1, diag2)
        if diag1 != diag2:
            yield min(diag1, diag2) if row_num % 2 else max(diag1, diag2)

print list(matrix_diagonals(4))


You could even simplify the arithmetic a bit:

def matrix_diagonals(n):
    for row_num in xrange(n):
        diag1 = (n + 1) * row_num + 1
        diag2 = (n - 1) * row_num + n
        yield max(diag1, diag2) if row_num % 2 else min(diag1, diag2)
        if diag1 != diag2:
            yield min(diag1, diag2) if row_num % 2 else max(diag1, diag2)


Since you are using Python 2, xrange() would be more efficient than range().