python - 使用 TxK 列索引数组从 TxN numpy 数组中选取 TxK numpy 数组

Question

这是一个间接索引问题。

可以通过列表理解来解决。

问题是是否或者如何在 numpy 中解决它，

什么时候 data.shape 是 (T,N)和 c.shape 是 (T,K)

并且c的每个元素都是一个介于0和N-1之间的int，即:c 的每个元素旨在引用 data 中的列号。

目标是获取出，其中

out.shape = (T,K)

对于 0..(T-1) 中的每个 i

行out[i] = [ data[i, c[i,0]] , ... , data[i, c[i,K-1]] ]

具体例子:

data = np.array([\
       [ 0,  1,  2],\
       [ 3,  4,  5],\
       [ 6,  7,  8],\
       [ 9, 10, 11],\
       [12, 13, 14]])

c = np.array([
      [0, 2],\
      [1, 2],\
      [0, 0],\       
      [1, 1],\       
      [2, 2]])

out should be out = [[0, 2], [4, 5], [6, 6], [10, 10], [14, 14]]

out 的第一行是 [0,2]，因为选择的列是由 c 的第 0 行给出的，它们是 0 和 2，第 0 和 2 列的 data[0] 是 0 和 2。

out的第二行是[4,5]，因为选择的列是由c的第1行给出的，它们是1和2，而第1和2列的data[1]是4和5。

Numpy 花式索引似乎并没有以明显的方式解决这个问题，因为使用 c 索引数据(例如 data[c]、np.take(data,c,axis=1 ) ) 总是生成一个 3 维数组。

列表理解可以解决这个问题:

out = [ [data[rowidx,i1],data[rowidx,i2]] for (rowidx, (i1,i2)) in enumerate(c) ]

如果 K 是 2，我想这还算可以。如果 K 是可变的，这不太好。

必须为每个值 K 重写列表推导式，因为它会展开 c 的每一行从 data 中选取的列。这也违反了 DRY。

有没有完全基于numpy的解决方案？

Answer 1

您可以使用 np.choose 避免循环:

In [1]: %cpaste
Pasting code; enter '--' alone on the line to stop or use Ctrl-D.

data = np.array([\
       [ 0,  1,  2],\
       [ 3,  4,  5],\
       [ 6,  7,  8],\
       [ 9, 10, 11],\
       [12, 13, 14]])

c = np.array([
      [0, 2],\
      [1, 2],\
      [0, 0],\
      [1, 1],\
      [2, 2]])
--

In [2]: np.choose(c, data.T[:,:,np.newaxis])
Out[2]: 
array([[ 0,  2],
       [ 4,  5],
       [ 6,  6],
       [10, 10],
       [14, 14]])