时间戳

首页 > 解决方案 > 基于第二个 Numpy 数组的更快的 Numpy 数组元素分组方法

python - 基于第二个 Numpy 数组的更快的 Numpy 数组元素分组方法

问题描述

有 2 个 NumPy 数组groups和selectors，其中

  1. selectors是一个包含需要分组的整数的数组
import numpy as np
np.random.seed(0)

selectors = np.random.randint(0, 300, 5)
# [172  47 117 192 251]
  1. groups是一个结构化数组，包含组 (str) 的第一个索引 (int)
# Generate groups `a` to `t` and their first index
start = ord('a')
groups = []
for i in range(20):
    e = (i*i, chr(start+i))
    groups.append(e)
groups = np.array(groups, dtype=[('index', np.uint32), ('selector', '|U1')])
groups = np.sort(groups, order='index')

# [(  0, 'a') (  1, 'b') (  4, 'c') (  9, 'd') ( 16, 'e') ( 25, 'f')
#  ( 36, 'g') ( 49, 'h') ( 64, 'i') ( 81, 'j') (100, 'k') (121, 'l')
#  (144, 'm') (169, 'n') (196, 'o') (225, 'p') (256, 'q') (289, 'r')
#  (324, 's') (361, 't')]

给定这些示例数组，分组后所需的结果将是np.ndarrays/lists的字典

{
    "g": [47] ,
    "k": [117],
    "n": [172, 192],
    "p": [251]
}

有没有更快的方法在 Numpy 中执行此分组而不是嵌套 2 个循环，如下所示？这对于使用100-1000 行的selectors数组的 10-1 亿行的大型数组很有用。groups

使用嵌套循环

results = {}
for s in selectors:
    for i in range(len(groups)-1):
        if s >= groups[i][0] and s < groups[i+1][0]:
            j = i
            break
    else:
        j = i + 1

    try:
        results[groups[j][1]].append(s)
    except KeyError:
        results[groups[j][1]] = [s]
print(results)
# {'n': [172, 192], 'g': [47], 'k': [117], 'p': [251]}

标签： pythonarraysalgorithmnumpystructured-array

解决方案

如果您在每个选择器上使用二进制搜索，您实际上是在将例程的时间O(len(groups) * len(selectors))从O(log2(len(groups)) * len(selectors))

该模块上的Python 文档bisect解释了如何使用它来查找小于或等于指定值的最右侧元素。

推荐阅读