如何從列索引矩陣設置矩陣的單元格

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]