编写一个函数，在Python语言中将数组(NxN)的对角元素作为数组返回

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(ai)导致错误显示：#TypeError:对于下面的代码，元组索引必须是整数，而不是字符串

Answer 1

如果你知道你的矩阵将是一个正方形的矩阵，你不需要单独接受行和列的输入。你只需要取值n，你的对角线函数看起来就像：

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

Answer 2

我确实同意@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]