Numpy根据另一个数组的值汇总一个数组

Question

我正在寻找一种向量化的方法来完成以下任务：

假设我有一组 x 和 y 的值。需要注意的是，x 的值不一定都是整数，而且可能是负数：

import numpy as np
x = np.array([-1,-1,-1,3,2,2,2,5,4,4], dtype=float)
y = np.array([0,1,0,1,0,1,0,1,0,1])

我想根据 x 数组中排序后的唯一值来对 y 数组进行分组，并总结每个 y 类的计数。所以上面的例子看起来应该是这样的：

array([[ 2.,  1.],
      [ 2.,  1.],
      [ 0.,  1.],
      [ 1.,  1.],
      [ 0.,  1.]])

第一列表示每个唯一的 x 值对应的 '0' 值的计数，第二列表示每个唯一的 x 值对应的 '1' 值的计数。

我目前的实现方式是这样的：

x_sorted, y_sorted = x[x.argsort()], y[x.argsort()]

def collapse(x_sorted, y_sorted):
     uniq_ids = np.unique(x_sorted, return_index=True)[1]
     y_collapsed = np.zeros((len(uniq_ids), 2))
     x_collapsed = x_sorted[uniq_ids]
     for idx, y in enumerate(np.split(y_sorted, uniq_ids[1:])):
          y_collapsed[idx,0] = (y == 0).sum()
          y_collapsed[idx,1] = (y == 1).sum()
     return (x_collapsed, y_collapsed)

collapse(x_sorted, y_sorted)
(array([-1, 2, 3, 4, 5]),
 array([[ 2.,  1.],
       [ 2.,  1.],
       [ 0.,  1.],
       [ 1.,  1.],
       [ 0.,  1.]]))

不过，这样的做法似乎不太符合 numpy 的精神，我希望能找到一种向量化的方法来完成这个操作。我想在不使用 pandas 的情况下做到这一点。我知道那个库有一个非常方便的分组操作。

Answer 1

这里还有一个解决方案：

y = y[np.argsort(x)]

b = np.bincount(x)
b = b[b!=0]

ans = np.array([[i.shape[0], i.sum()] for i in np.split(y, np.cumsum(b))[:-1]])

ans[:,0] -= ans[:,1]

print(ans)
#array([[2, 1],
#   [2, 1],
#   [0, 1],
#   [1, 1],
#   [0, 1]], dtype=int64)

时间安排：

 @seikichi solution:
 10000 loops, best of 3: 37.2 µs per loop

 @acushner solution:
 10000 loops, best of 3: 65.4 µs per loop

 @SaulloCastro solution:
 10000 loops, best of 3: 154 µs per loop

Answer 2

我还没测试过这个，但我觉得应该可以用。基本上，我就是根据x的值来获取y中的值。

uniques = list(set(x))
uniques.sort()
lu = len(uniques)
res = np.zeros(lu * 2).reshape(lu, 2)
for i, v in enumerate(uniques):
    cur = y[x == v]
    s = cur.sum()
    res[i, 0] = len(cur) - s
    res[i, 1] = s

另一种方法是使用numpy MaskedArrays。

Answer 3

np.unique 和 np.bincount 是你在这里的好帮手。下面的代码可以处理任何类型的输入，不一定是小的连续整数：

>>> x = np.array([1, 1, 1, 3, 2, 2, 2, 5, 4, 4])
>>> y = np.array([0, 1, 2, 2, 0, 1, 0, 2, 2, 1])
>>> 
>>> x_unq, x_idx = np.unique(x, return_inverse=True)
>>> y_unq, y_idx = np.unique(y, return_inverse=True)
>>> 
>>> np.column_stack(np.bincount(x_idx, y_idx == j) for j in range(len(y_unq)))
array([[ 1.,  1.,  1.],
       [ 2.,  1.,  0.],
       [ 0.,  0.,  1.],
       [ 0.,  1.,  1.],
       [ 0.,  0.,  1.]])

你还可以提取行和列的标签：

>>> x_unq
array([1, 2, 3, 4, 5])
>>> y_unq
array([0, 1, 2])

Answer 4

下面的代码怎么样？（使用 numpy.bincount 和 numpy.concatenate）

>>> import numpy as np
>>> x = np.array([1,1,1,3,2,2,2,5,4,4])
>>> y = np.array([0,1,0,1,0,1,0,1,0,1])
>>> xmax = x.max()
>>> numpy.concatenate([[numpy.bincount(x[y == v], minlength=xmax + 1)] for v in sorted(set(y))], axis=0)[:, 1:].T
array([[2, 1],
       [2, 1],
       [0, 1],
       [1, 1],
       [0, 1]])

---

更新：谢谢 @askewchan！

>>> import numpy as np
>>> x = np.array([1.1, 1.1, 1.1, 3.3, 2.2, 2.2, 2.2, 5.5, 4.4, 4.4])
>>> y = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
>>> r = np.r_[np.unique(x), np.inf]
>>> np.array([np.histogram(x[y == v], r)[0] for v in sorted(set(y))]).T
array([[2, 1],
       [2, 1],
       [0, 1],
       [1, 1],
       [0, 1]])

Answer 5

因为 x 是 float 类型的。我会这样做：

In [136]:

np.array([(x[y==0]==np.unique(x)[..., np.newaxis]).sum(axis=1),
          (x[y==1]==np.unique(x)[..., np.newaxis]).sum(axis=1)]).T
Out[136]:
array([[2, 1],
       [2, 1],
       [0, 1],
       [1, 1],
       [0, 1]])

速度：

In [152]:

%%timeit
ux=np.unique(x)[..., np.newaxis]
np.array([(x[y==0]==ux).sum(axis=1),
          (x[y==1]==ux).sum(axis=1)]).T
10000 loops, best of 3: 92.7 µs per loop

解决方案 @seikichi

In [151]:

%%timeit
>>> x = np.array([1.1, 1.1, 1.1, 3.3, 2.2, 2.2, 2.2, 5.5, 4.4, 4.4])
>>> y = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
>>> r = np.r_[np.unique(x), np.inf]
>>> np.concatenate([[np.histogram(x[y == v], r)[0]] for v in sorted(set(y))]).T
1000 loops, best of 3: 388 µs per loop

对于更一般的情况，当 y 不仅仅是 {0,1} 时，正如 @askewchan 指出的那样：

In [155]:

%%timeit
ux=np.unique(x)[..., np.newaxis]
uy=np.unique(y)
np.asanyarray([(x[y==v]==ux).sum(axis=1) for v in uy]).T
10000 loops, best of 3: 116 µs per loop

为了更进一步解释广播的概念，可以看这个例子：

In [5]:

np.unique(a)
Out[5]:
array([ 0. ,  0.2,  0.4,  0.5,  0.6,  1.1,  1.5,  1.6,  1.7,  2. ])
In [8]:

np.unique(a)[...,np.newaxis] #what [..., np.newaxis] will do:
Out[8]:
array([[ 0. ],
       [ 0.2],
       [ 0.4],
       [ 0.5],
       [ 0.6],
       [ 1.1],
       [ 1.5],
       [ 1.6],
       [ 1.7],
       [ 2. ]])
In [10]:

(a==np.unique(a)[...,np.newaxis]).astype('int') #then we can boardcast (converted to int for readability)
Out[10]:
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
       [1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0]])
In [11]:

(a==np.unique(a)[...,np.newaxis]).sum(axis=1) #getting the count of unique value becomes summing among the 2nd axis
Out[11]:
array([1, 3, 1, 1, 2, 1, 1, 1, 1, 3])