时间戳

首页 > 解决方案 > 沿列展平numpy数组，顺序为：下三角形，对角线，上三角形

python - 沿列展平numpy数组，顺序为：下三角形，对角线，上三角形

问题描述

这是我试图解决的一个问题。假设我们有一个方阵：

In [10]: arr
Out[10]: 
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12],
       [13, 14, 15, 16]])

我想要的是按特定顺序展平这个数组：首先我想沿轴 0 展平下三角形，然后选择对角线，最后再次沿轴 0 展平上三角形，这最终会将展平数组设为：

#              | lower triangle     |diag.elements| upper triangle  |  
res = np.array([5, 9, 13, 10, 14, 15, 1, 6, 11, 16, 2, 3, 7, 4, 8, 12])

到目前为止，这是我的部分解决方案，但还没有给出预期的结果。

In [16]: arr[np.tril(arr, k=-1) != 0]
Out[16]: array([ 5,  9, 10, 13, 14, 15])   # not correct!

In [17]: np.diag(arr)
Out[17]: array([ 1,  6, 11, 16])

In [18]: arr[np.triu(arr, k=1) != 0]
Out[18]: array([ 2,  3,  4,  7,  8, 12])  # not correct!

最后，连接这 3 个中间结果。如何正确索引以获得期望的结果？或者，还有其他方法可以解决这个问题吗？

标签： pythonnumpymultidimensional-arrayflatten

解决方案

这是一个基于masking和concatenating/stacking-

In [50]: r = np.arange(len(arr))

In [51]: mask = r[:,None]<r

In [54]: np.concatenate((arr.T[mask],np.diag(arr),arr.T[mask.T]))
Out[54]: array([ 5,  9, 13, 10, 14, 15,  1,  6, 11, 16,  2,  3,  7,  4,  8, 12])

另一个完全基于masking-

n = len(arr)
r = np.arange(n)
mask = r[:,None]<r
diag_mask = r[:,None]==r
comp_mask = np.vstack((mask[None],diag_mask[None],mask.T[None]))
out = np.broadcast_to(arr.T,(3,n,n))[comp_mask]

推荐阅读