使用行索引为 numpy 数组赋值

Question

假设我有两个数组，a=np.array([0,0,1,1,1,2]), b=np.array([1,2,4,2,6,5])。 a 中的元素表示应分配 b 的行索引。如果同一行中有多个元素，则应按顺序分配值。 所以结果是一个二维数组 c:

c = np.zeros((3, 4))
counts = {k:0 for k in range(3)}
for i in range(a.shape[0]):
    c[a[i], counts[a[i]]]=b[i]
    counts[a[i]]+=1
print(c)

有没有办法在 numpy 中使用一些奇特的索引方法来更快地获得这样的结果（没有 for 循环）以防这些数组很大。

Answer 1

我必须 运行 你的代码才能真正看到它产生了什么。我脑子里的东西是有限度的'run'。

In [230]: c                                                                                            
Out[230]: 
array([[1., 2., 0., 0.],
       [4., 2., 6., 0.],
       [5., 0., 0., 0.]])
In [231]: counts                                                                                       
Out[231]: {0: 2, 1: 3, 2: 1}

省略此信息可能会延迟可能的答案。 'vectorization' 需要从整个数组的角度思考，如果我可以将结果可视化并寻找模式，这将是最简单的。

这看起来像一个 padding 问题。

In [260]: u, c = np.unique(a, return_counts=True)                                                      
In [261]: u                                                                                            
Out[261]: array([0, 1, 2])
In [262]: c                                                                                            
Out[262]: array([2, 3, 1])      # cf with counts

根据之前的填充问题，我可以构造一个掩码：

In [263]: mask = np.arange(4)<c[:,None]                                                                
In [264]: mask                                                                                         
Out[264]: 
array([[ True,  True, False, False],
       [ True,  True,  True, False],
       [ True, False, False, False]])

并使用它将 b 值分配给 c:

In [265]: c = np.zeros((3,4),int)                                                                      
In [266]: c[mask] = b                                                                                  
In [267]: c                                                                                            
Out[267]: 
array([[1, 2, 0, 0],
       [4, 2, 6, 0],
       [5, 0, 0, 0]])

由于 a 已经排序，我们可能比 unique 更快地获得计数。如果 a 某些行没有任何值，它也会有问题。