快速选择范围之间的元素百分比

Question

给定预定义的范围、百分比列表和一些数据，我需要从位于每个范围之间的元素中随机选择一个百分比的 ID。

下面的代码显示了我是如何做到的，for 块目前是瓶颈。我确信这可能会通过一些矢量化来加快速度，但我不知道如何。

import numpy as np
import itertools

# Generate some random data
N = 1000
aa = np.random.uniform(12., 20., N)
# Define edges/ranges.
edges = np.array([16.67666667, 16.77721569, 16.87776471, 16.97831373,
                  17.07886275, 17.17941176, 17.27996078, 17.3805098,
                  17.48105882, 17.58160784, 17.68215686, 17.78270588,
                  17.8832549, 17.98380392, 18.08435294, 18.18490196,
                  18.28545098, 18.386])
# Percentage of elements in 'aa' that will be kept for each 'edges' range.
perc = np.random.uniform(0., 1., len(edges) - 1)

# Locate indexes of 'aa' elements within each 'edges' range.
c_indx = np.searchsorted(edges, aa, side='left')

# THIS IS THE BOTTLENECK
cc = []
# For each defined percentage value (one per edge range).
for i, p in enumerate(perc):
    # Locate IDs of lements within each range. Use 'i + 1' since the first
    # edge range (ie: those elements with c_indx=0) are discarded.
    idxs = np.where(c_indx == i + 1)[0]
    # Shuffle IDs within this edge range (in place)
    np.random.shuffle(idxs)
    # Estimate the number of elements from 'aa' to keep for
    # this range, given the fixed percentage 'p'.
    d = int(round(idxs.size * p, 0))
    # Store the correct percentage of random IDs from 'aa' for this range.
    cc.append(idxs[:d])

# Final list (flattened)
cc = list(itertools.chain.from_iterable(cc))

Answer 1

我们可以通过在进入循环之前计算c_indx 的排序索引来计算idxs，从而减少循环内的工作量。

因此，一种解决方案是 -

sidx = c_indx.argsort()
sc = c_indx[sidx]
idx = np.flatnonzero(sc[1:] != sc[:-1])+1
for i, p in enumerate(perc):
    idxs = sidx[idx[i]:idx[i+1]]
    # ..rest of the code stays the same

Answer 2

您已经从 searchsorted 调用中的 np.where 获得了信息。为什么不使用字典？

mydict = {}

for i,j in enumerate(c_indx):
    mydict.setdefault(j,[]).append(i)

这部分代码对我来说需要 57.2us，而不是 300。我们也可以将数据保存在列表中以加快速度。

cc = []

for i, p in enumerate(perc):
    # Locate IDs of lements within each range. Use 'i + 1' since the first
    # edge range (ie: those elements with c_indx=0) are discarded.
    idxs = mydict[i + 1]
    # Shuffle IDs within this edge range (in place)
    np.random.shuffle(idxs)
    # Estimate the number of elements from 'aa' to keep for this range, given the fixed percentage 'p'.
    d = int(round(len(idxs) * p, 0))
    # Store the correct percentage of random IDs from 'aa' for this range.
    cc.append(idxs[:d])