NumPy：构造矩阵对角线上的正方形/扩展对角线矩阵

Question

假设你有两个数组:

index = [1, 2, 3]
counts = [2, 3, 2]

或一个单一的数组

arr = [1, 1, 2, 2, 2, 3, 3]

我该如何高效构建这个矩阵

[
  [1, 1, 0, 0, 0, 0, 0],
  [1, 1, 0, 0, 0, 0, 0],
  [0, 0, 2, 2, 2, 0, 0],
  [0, 0, 2, 2, 2, 0, 0],
  [0, 0, 2, 2, 2, 0, 0],
  [0, 0, 0, 0, 0, 3, 3],
  [0, 0, 0, 0, 0, 3, 3]
]

使用NumPy吗？

我知道

square = np.zeros((7, 7))
np.fill_diagnol(square, arr) # see arr above

产生

[
  [1, 0, 0, 0, 0, 0, 0],
  [0, 1, 0, 0, 0, 0, 0],
  [0, 0, 2, 0, 0, 0, 0],
  [0, 0, 0, 2, 0, 0, 0],
  [0, 0, 0, 0, 2, 0, 0],
  [0, 0, 0, 0, 0, 3, 0],
  [0, 0, 0, 0, 0, 0, 3]
]

如何通过n来“扩展”对角线，其中n是由index[I]指定的值的counts[index-1]

tmp = np.array((arr * N)).reshape((len(arr), len(arr)) 
np.floor( (tmp + tmp.T) / 2 ) # <-- this is closer


array([[1., 1., 1., 1., 1., 2., 2.],
       [1., 1., 1., 1., 1., 2., 2.],
       [1., 1., 2., 2., 2., 2., 2.],
       [1., 1., 2., 2., 2., 2., 2.],
       [1., 1., 2., 2., 2., 2., 2.],
       [2., 2., 2., 2., 2., 3., 3.],
       [2., 2., 2., 2., 2., 3., 3.]])

这可以得到我想要的，但可能扩展性不太好？

riffled = list(zip(index, counts))
riffled
# [(1, 2), (2, 3), (3, 2)]

a = np.zeros((len(arr), len(arr))) # 7, 7 square
last = 0 # <-- keep track of current sub square
for i, c in riffled:
    a[last:last+c, last:last+c] = np.ones((c, c)) * i 
    last += c # <-- shift square

产量

array([[1., 1., 0., 0., 0., 0., 0.],
       [1., 1., 0., 0., 0., 0., 0.],
       [0., 0., 2., 2., 2., 0., 0.],
       [0., 0., 2., 2., 2., 0., 0.],
       [0., 0., 2., 2., 2., 0., 0.],
       [0., 0., 0., 0., 0., 3., 3.],
       [0., 0., 0., 0., 0., 3., 3.]])

- SumNeuron

---

np.equal.outer(arr, arr) * arr - user3483203

@user3483203 也可以使用！谢谢。 - SumNeuron

Answer 1

你可以使用scipy.linalg.block_diag来实现这一点：

import numpy as np
import scipy.linalg as linalg


a = 1*np.ones((2,2))
b = 2*np.ones((3,3))
c = 3*np.ones((2,2))

superBlock = linalg.block_diag(a,b,c)

print(superBlock)

#returns
#[[1. 1. 0. 0. 0. 0. 0.]
# [1. 1. 0. 0. 0. 0. 0.]
# [0. 0. 2. 2. 2. 0. 0.]
# [0. 0. 2. 2. 2. 0. 0.]
# [0. 0. 2. 2. 2. 0. 0.]
# [0. 0. 0. 0. 0. 3. 3.]
# [0. 0. 0. 0. 0. 3. 3.]]

如果你想从一个值列表和一个数量列表中获取到某个元素，可以这样做：

values = [1,2,3]
counts = [2,3,2]

mats = []
for v,c in zip(values,counts):
    thisMatrix = v*np.ones((c,c))
    mats.append( thisMatrix )


superBlock = linalg.block_diag(*mats)


print(superBlock)

- Ulises Bussi

---

所以 mats = [np.full((c, c), v) for v, c in zip(values, counts)]？ - Mad Physicist

这是一个很好的方法。 - Quang Hoang

如果用户在创建步骤之后只需要使用密集矩阵，则这非常完美。然而，如果所有数据都可以像示例一样被稀疏地描述，那么使用稀疏版本也可能会有一些好处：scipy.sparse.block_diag https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.block_diag.html#scipy.sparse.block_diag - Mike Holcomb

@MikeHolcomb想提出修改意见吗？这将有助于得到更好的答案（或者您可以自己回答，我会点赞）。 - Ulises Bussi

Answer 2

这里有一个通用的解决方案。

从索引/计数开始：

index = [1, 2, 1]
counts = [2, 3, 2]

arr = np.repeat(index, counts)
arr2 = np.repeat(range(len(index)), counts)
np.where(arr2 == arr2[:, None], arr, 0)

输出：

array([[1, 1, 0, 0, 0, 0, 0],
       [1, 1, 0, 0, 0, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 0, 0, 0, 1, 1],
       [0, 0, 0, 0, 0, 1, 1]])

从数组版本开始：

arr = np.array([1, 1, 2, 2, 2, 1, 2])

arr2 = np.cumsum(np.diff(arr,prepend=np.nan) != 0)
np.where(arr2 == arr2[:, None], arr, 0)

输出：

array([[1, 1, 0, 0, 0, 0, 0],
       [1, 1, 0, 0, 0, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 0, 0, 0, 1, 0],
       [0, 0, 0, 0, 0, 0, 2]])

- mozway

Answer 3

尝试使用广播：

idx = np.repeat(np.arange(len(counts)), counts)
np.where(idx==idx[:,None], arr, 0)
# or
# arr * (idx==idx[:,None])

输出；

array([[1, 1, 0, 0, 0, 0, 0],
       [1, 1, 0, 0, 0, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 2, 2, 2, 0, 0],
       [0, 0, 0, 0, 0, 3, 3],
       [0, 0, 0, 0, 0, 3, 3]])

- Quang Hoang

---

没问题！请注意，这里提到的 arr 是 Python 的列表，所以为了让它起作用，需要通过 np.array(arr) 将其转换为 numpy 数组。 - SumNeuron

@SumNeuron。你可以使用np.reshape(arr, (-1, 1)) == arr来避免显式转换。不过显式转换更便宜。 - Mad Physicist

请注意，仅当arr中的块不同（例如，arr = np.array([1, 1, 2, 2, 2, 1, 1])将无法工作）。 - mozway

@mozway请查看更新后的答案。我们不仅仅是不使用npindex作为参考。 - Quang Hoang

@QuangHoang 我刚刚发布了类似的内容 ;) （抱歉回复晚了，我不在电脑前） - mozway