如何从列索引矩阵设置矩阵的单元格

Question

我想从位置列表和内核中心列表构建内核。内核应该是两个离每个位置最近的中心的指标。

> x = np.array([0.1, .49, 1.9, ]).reshape((3,1)) # Positions 
> c = np.array([-2., 0.1, 0.2, 0.4, 0.5, 2.]) # centers 
print x 
print c 

[[ 0.1 ] 
[ 0.49] 
[ 1.9 ]] 
[-2. 0.1 0.2 0.4 0.5 2. ] 

我想获得的是：

array([[ 0, 1, 1, 0, 0, 0], # Index 1,2 closest to 0.1 
     [ 0, 0, 0, 1, 1, 0], # Index 3,4 closest to 0.49 
     [ 0, 0, 0, 0, 1, 1]]) # Index 4,5 closest to 1.9 

我可以得到：

> dist = np.abs(x-c) 
array([[ 2.1 , 0. , 0.1 , 0.3 , 0.4 , 1.9 ], 
     [ 2.49, 0.39, 0.29, 0.09, 0.01, 1.51], 
     [ 3.9 , 1.8 , 1.7 , 1.5 , 1.4 , 0.1 ]]) 

和：

> np.argsort(dist, axis=1)[:,:2] 
array([[1, 2], 
     [4, 3], 
     [5, 4]]) 

在这里，我有一个矩阵列索引，但我但看不到如何使用它们来设置这些值另一个矩阵中的列（使用高效的numpy操作）。

idx = np.argsort(dist, axis=1)[:,:2] 
z = np.zeros(dist.shape) 
z[idx]=1 # NOPE 
z[idx,:]=1 # NOPE 
z[:,idx]=1 # NOPE 

Answer 1

一种方法是初始化为零数组，然后用advanced-indexing指数 -

out = np.zeros(dist.shape,dtype=int) 
out[np.arange(idx.shape[0])[:,None],idx] = 1 

另外，我们可以玩的尺寸扩展使用broadcasting，并拿出一班轮 -

out = (idx[...,None] == np.arange(dist.shape[1])).any(1).astype(int) 

出于性能考虑，我建议使用np.argpartition得到这些指数 -

idx = np.argpartition(dist, 2, axis=1)[:,:2]