python 实现方阵的对角线遍历

python 实现方阵的对角线遍历

---

任务描述

对一个方阵矩阵，实现平行于主对角线方向的对角线元素遍历。

---

• 从矩阵索引入手：

[[ 1  2  3  4  5]
 [ 6  7  8  9 10]
 [11 12 13 14 15]
 [16 17 18 19 20]
 [21 22 23 24 25]]

• 上三角的索引遍历：

0 0
1 1
2 2
3 3
4 4

0 1
1 2
2 3
3 4

0 2
1 3
2 4

0 3
1 4

0 4

• 下三角的索引遍历：

1 0
2 1
3 2
4 3

2 0
3 1
4 2

3 0
4 1

4 0

---

代码

import numpy as np

A = np.arange(25)+1
A = np.mat(A.reshape([5, 5]))
print(A)

"""
[[ 1  2  3  4  5]
 [ 6  7  8  9 10]
 [11 12 13 14 15]
 [16 17 18 19 20]
 [21 22 23 24 25]]
"""

Num_element = A.shape[0]
c = int((Num_element-1)/2)
# print(c)


R = np.zeros_like(A)
# print(R)

for j in range(Num_element):
	print()
	i = 0
	# print(i, j)
	while np.max([i, j])<Num_element:
		print(i, j)
		if np.abs(i-j)%2==0:
			R[i, j] = A[c-int((j-i)/2), c+int((j-i)/2)]
		else:
			R[i, j] = (A[c-int((j-i-1)/2), c+int((j-i+1)/2)]+A[c-int((j-i+1)/2), c+int((j-i-1)/2)])/2

		i=i+1
		j=j+1

# print(R)

for k in range(1, Num_element):
	print()
	i = 0
	# print(i, j)
	while np.max([k, i])<Num_element:
		print(k, i)
		if np.abs(k-i)%2==0:
			R[k, i] = A[c-int((i-k)/2), c+int((i-k)/2)]
		else:
			R[k, i] = (A[c-int((i-k-1)/2), c+int((i-k+1)/2)]+A[c-int((i-k+1)/2), c+int((i-k-1)/2)])/2

		k=k+1
		i=i+1

print(R)

上述代码中对于每条对角线的所有元素执行相同的赋值操作。

考虑将其中重复的部分封装成函数：

def diag_opreation(k, i, Num_element, R, A):
	c = int((Num_element-1)/2)
	while np.max([k, i])<Num_element:
			print(k, i)
			if np.abs(k-i)%2==0:
				R[k, i] = A[c-int((i-k)/2), c+int((i-k)/2)]
			else:
				R[k, i] = (A[c-int((i-k-1)/2), c+int((i-k+1)/2)]+A[c-int((i-k+1)/2), c+int((i-k-1)/2)])/2
	
			k=k+1
			i=i+1
	return R

则代码变为：

for j in range(Num_element):
	print()
	i = 0
	# print(i, j)
	R = diag_opreation(i, j, Num_element, R, A)

# print(R)

for k in range(1, Num_element):
	print()
	i = 0
	# print(i, j)
	R = diag_opreation(k, i, Num_element, R, A)

print(R)

输出结果为：

[[13 11  9  7  5]
 [15 13 11  9  7]
 [17 15 13 11  9]
 [19 17 15 13 11]
 [21 19 17 15 13]]