这是一个间接索引问题。
可以通过列表理解来解决。
问题是是否或者如何在 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
的解决方案?
最佳答案
您可以使用 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]])
关于python - 使用 TxK 列索引数组从 TxN numpy 数组中选取 TxK numpy 数组,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/26222835/