A natural number N is introduced. Using the list comprehension, form a two-dimensional list of size N x N, consisting of zeros and ones along the main diagonal. (The main diagonal is the elements running diagonally from the upper left corner of the matrix to its lower right corner). Output the result as a table of numbers as shown in the example below.

def main_diag(a):
  return [[1 if i == j else 0 for i in range(a)] for j in range(a)]
main_diag(5)