Напишите функцию для возврата диагональных элементов массива (NxN) в виде массива в Python

Question

import numpy as np
#getting number of rows and columns for the arrat
nr=input("enter rows")
nc=input("enter columns")
print("Please enter same rows and columns")
n=nr*nc
ar1=[]
#checking for square matrix
if(nr==nc):
#loop to append elements into the empty-list ar1
  for i in range(n):
     ele=input("enter elements")
     ar1.append(ele)
#getting the number of rows of array1
array1=np.array(ar1).shape
#function to get the diagonal elements of array
def diagonal(a):
  global n
  n=str(n)
  for i in n:
     for j in n:
       if i==j:

          newarr=np.array(a[i][j])

  #print(newarr) 
diagonal(array1)

newarr = np.array (a [i] [j]) приводит к ошибке, показывающей: #TypeError: индексы кортежа должны быть целыми числами, а не str для приведенного ниже кода

Answer 1

Я согласен с @Gautam, кстати, вот решение, которое использует встроенную функцию diagonal:

lst = []
size = int(input('Enter size of sqaure matrix'))
for i in range(1, (size**2)+1):
    lst.append(int(input('Enter ' + str(i) +'th element')))
lst

Enter size of sqaure matrix2
Enter 1th element5
Enter 2th element6
Enter 3th element8
Enter 4th element7
[5, 6, 8, 7]

arr = np.array(lst).reshape(size,size)
arr.diagonal()

[5, 7]

Answer 2

Если вы знаете, что ваша матрица будет квадратной, вам не нужно отдельно вводить строки и столбцы.Вы можете просто принять значение n.Ваша диагональная функция будет выглядеть следующим образом:

newa=[]
def diag(arr):
    for i,a in enumerate(arr):
        newa[i]=arr[i,i]