计算排序数组子集中唯一值的数量

Question

我有两个 numpy 数组，users 和 dat。对于 users 中的每个用户，我需要在 dat 中找到与该用户相关的数据并计算唯一值的数量。我需要处理 len(users)=200000 和 len(dat)=2800000 的情况。目前我没有利用 dat 被排序的事实，这使得该方法非常慢。我该怎么做？

dat 中的值 'other' 仅表明其他值也将出现在结构化数组中。

import numpy as np

users = np.array([111, 222, 333])
info = np.zeros(len(users))
dt = [('id', np.int32), ('group', np.int16), ('other', np.float)]
dat = np.array([(111, 1, 0.0), (111, 3, 0.0), (111, 2, 0.0), (111, 1, 0.0),
               (222, 1, 0.0), (222, 1, 0.0), (222, 4, 0.0),
               (333, 2, 0.0), (333, 1, 0.0), (333, 2, 0.0)],
               dtype=dt)

for i, u in enumerate(users):
    u_dat = dat[np.in1d(dat['id'], u)]
    uniq = set(u_dat['group'])
    info[i] = int(len(uniq))

print info

Answer 1

如果你想从 numpy 的矢量化中获益，如果你能事先从 dat 中删除所有重复项，那将大有帮助。然后，您可以通过两次调用 searchsorted:

找到值的第一次出现和最后一次出现

dat_unq = np.unique(dat)
first = dat_unq['id'].searchsorted(users, side='left')
last =  dat_unq['id'].searchsorted(users, side='right')
info = last - first

只有当您要在 dat 中搜索大量条目时，这才有用。如果它是一个较小的分数，您仍然可以使用对 searchsorted 的两次调用来确定要在 unique 上调用哪些切片：

info = np.empty_like(users, dtype=np.intp)
first = dat['id'].searchsorted(users, side='left')
last =  dat['id'].searchsorted(users, side='right')
for idx, (start, stop) in enumerate(zip(first, last)):
    info[idx] = len(np.unique(dat[start:stop]))

计算排序数组子集中唯一值的数量

Counting number of unique values in subset of sorted array

python

numpy

optimization

structured-array