python - 快速选择范围内一定百分比的元素

Question

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

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

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