python - 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 的情况下做到这一点。我知道图书馆有一个非常方便的 groupby 操作。

Answer 1

因为 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])