Print diagonal numbers of matrix in ascending order

Question

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)

Answer 1

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().

Answer 2

You can improve the create_matrix() function by using a list comprehension while reversing the inner list direction if a row number is odd:

def create_matrix(n):
 """
 Generates and returns a matrix according to the following rules:
 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):
 1 2 3 4
 8 7 6 5
 9 10 11 12
 16 15 14 13
 """
 return [[col + n * row + 1 for col in range(n)][::1 if row % 2 == 0 else -1]
 for row in range(n)]

^{(I think we can further improve the 1/-1 direction check)}

Note that the create_matrix(), according to it's name, should create a matrix and return it. Don't call other functions from it to not disrupt separation of concerns and modularity. Something like:

matrix = create_matrix(n)
print(get_sorted_diagonal_elements(matrix))

---

Also, here is an alternative implementation of the second part - it is though worse than your version in terms of space complexity and is not optimal at all, but I'll post if for educational reasons anyway. Here, I'm using the minimum heap to keep track of the minimum diagonal elements:

import heapq
def get_sorted_diagonal_elements(matrix, n):
 min_heap = []
 for row in range(n):
 heapq.heappush(min_heap, matrix[row][row])
 heapq.heappush(min_heap, matrix[row][n - row - 1])
 for _ in range(n * 2):
 yield heapq.heappop(min_heap)
n = 4
matrix = create_matrix(n)
for diagonal_element in get_sorted_diagonal_elements(matrix, n):
 print(diagonal_element)

Answer 3

Thing is, the diagonal of a matrix numbered from 1 to \$n^2\$ follows a simple pattern:

$$ \begin{equation} 1 \\ n+2 \\ 2n+3 \\ 3n+4 \\ \vdots \\ n^2 \end{equation} $$

i.e.

$$\{(i-1)n + i, \forall i \in [1; n]\}$$

But here we have some kind of zigzag pattern:

$$ \begin{equation} 1 \\ 2n - 1 \\ 2n + 3 \\ 4n - 3 \\ \vdots \\ n^2 \end{equation} $$

and we should also consider the second diagonal:

$$ \begin{equation} n \\ n + 2 \\ 3n - 2 \\ 3n + 4 \\ \vdots \\ n^2 - n + 1 \end{equation} $$

... Not much of a pattern until we combine them:

$$ \begin{align} 0n+1 && 1n-0 \\ 1n+2 && 2n-1 \\ 2n+3 && 3n-2 \\ 3n+4 && 4n-3 \\ \vdots && \vdots \\ (n-1)n+n && n \times n-(n-1) \end{align} $$

There, much easier to generate programatically:

def generate_matrix_diagonals(size):
 for i, j in enumerate(xrange(size), 1):
 diag1 = j * size + i
 diag2 = i * size - j
 if diag1 < diag2:
 yield diag1
 yield diag2
 elif diag1 == diag2:
 yield diag1
 else:
 yield diag2
 yield diag1

The advantage here is that we can iterate over this generator and have \$O(1)\$ space and \$O(n)\$ complexity:

for diag in generate_matrix_diagonals(n):
 print diag

Or, if you need to print it on a single line with comas inbetween, you can do so using \$O(n)\$ space:

print ', '.join(map(str, generate_matrix_diagonals(n)))