优化的功能将数组分成 bin

Question

我正在 python 中创建一个函数，以根据 x 数组将两个数组 x 和 y 分成预定数量的 bin。我已经制作了一种有效的算法，但它真的很慢。这是（显然）有效的代码：

def sepbin(x, y, classes_number=100, log_scale=True):
    if log_scale:
        if x[0]<=0:
            print 'Warning: zero value in array about to be log-scaled. Ignoring it.'
            x=x[1:]
            y=y[1:]
        bins=np.logspace(np.log(x[0]), np.log(x[-1]), classes_number+1, base=np.e)
    else:
        bins=np.linspace(x[0], x[-1], classes_number+1)
    ybins=[[] for i in range(classes_number)]
    xbins=[[] for i in range(classes_number)]
    for xx, yy in zip(x,y):
        i=0
        while i<classes_number:
            if ((xx>=bins[i]) and (xx<bins[i+1])):
                ybins[i].append(yy)
                xbins[i].append(xx)
                break
            elif (i==(classes_number-1)) and xx==bins[-1]:
                ybins[i].append(yy)
                xbins[i].append(xx)
                break
            else:
                i+=1
    xsm = np.array(map(np.mean, xbins))
    ysm = np.array(map(np.mean, ybins))
    return xsm, ysm

如您所见，我想为对数缩放和线性缩放的输出腾出空间，所以我不能假设线性间隔的 bin。我只是假设数据是按新月或降序组织的（但这很容易概括）。

显然，该代码运行良好，但由于我希望使用非常大的数据集（100000 多个元素），我认为应该对其进行优化。有什么方法可以使用numpy 或scipyhere 以使其更快？我在numpy! 中找不到分箱功能，这让我有点吃惊，所以也许我看起来不太对。

谢谢。

Answer 1

您可以使用 numpy 的索引来查找每个 bin 中的 x 和 y 的元素，而不是 O(nk) 双循环：

def sepbin2(x, y, classes_number=100, log_scale=True):
    if log_scale:
        if x[0]<=0:
            print 'Warning: zero value in array about to be log-scaled. Ignoring it.'
            x=x[1:]
            y=y[1:]
        bins=np.logspace(np.log(x[0]), np.log(x[-1]), classes_number+1, base=np.e)
    else:
        bins=np.linspace(x[0], x[-1], classes_number+1)
    # pre-allocate return values
    xsm = np.zeros(classes_number)
    ysm = np.zeros(classes_number)
    # find elements in each bin
    for i in range(classes_number):
        if i == classes_number - 1:
            sel = bins[i] <= x
        else:
            sel = (bins[i] <= x) & (x < bins[i+1])
        xsm[i] = np.mean(x[sel])
        ysm[i] = np.mean(y[sel])
    return xsm, ysm

在我的测试中，这个函数的输出似乎和你写的完全一样。我使用稍微不同的代码来选择最后一个 bin（我不确定你为什么写 xx==bins[-1]）。

sepbin2 快很多。使用 100K 数据点：

x = np.random.random((100000,))
x.sort()
y = np.random.random((100000,))

我们得到：

In [1]: %timeit sepbin(x, y, classes_number=100)
1 loops, best of 3: 5.21 s per loop
In [2]: %timeit sepbin2(x, y, classes_number=100)
100 loops, best of 3: 18.9 ms per loop