无循环的复杂索引

时间:2018-08-14 10:10:39

标签: python-3.x numpy indexing mean

我有一个具有Nx5尺寸的数组。最后3个维度是不同体积的x,y和z。然后将它们打包到其他两个维度中。例如:

[[0 0 x0 y0 z0],
[0 0 x1 y1 z1],
[0 0 x2 y2 z2],
[0 1 x3 y3 z3],
[0 1 x4 y4 z4],
[1 0 x5 y5 z5],
[1 0 x6 y6 z6],
[1 1 x7 y7 z7],
[1 1 x8 y8 z8],
[2 0 x9 y9 z9],
[2 0 x10 y10 z10],
[2 0 x11 y11 z11],
[2 1 x12 y12 z12]]

前两个维度的卷数每次都不同。我想计算每个维度中每个体积的x,y和z平均值。这应该导致如下结果:

[[0 0 xmean0 ymean0 zmean0],
[0 1 xmean1 ymean1 zmean1],
[1 0 xmean2 ymean2 zmean2],
[1 1 xmean3 ymean3 zmean3]]
[2 0 xmean4 ymean4 zmean4],
[2 1 xmean5 ymean5 zmean5]]

换句话说,对于第一个to元素的每个组合,它应该具有均值。我不能为此使用循环,只能使用numpy和/或tensorflow。

1 个答案:

答案 0 :(得分:1)

我们假定输入数组为a

方法1:使用bincount-

unq_comb,ids, w = np.unique(a[:,:2], axis=0, return_inverse=1, return_counts=1)
out = np.empty((len(unq_comb),5))
out[:,:2] = unq_comb
for i in [2,3,4]:
    out[:,i] = np.bincount(ids, a[:,i])/w

方法2:使用sorting-

sidx = np.lexsort(a[:,:2].T)
b = a[sidx]
idx = np.flatnonzero(np.r_[True,(b[:-1,:2] != b[1:,:2]).any(1),True])
w = np.diff(idx)[:,None].astype(float)
out = np.empty((len(unq_comb),5))
out[:,:2] = b[idx[:-1],:2]
out[:,2:] = np.add.reduceat(b[:,2:], idx[:-1], axis=0)/w