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

时间:2015-08-30 20:37:37

标签: python numpy optimization structured-array

我有两个numpy数组,usersdat。对于users中的每个用户,我需要在dat中找到与用户相关的数据,并计算唯一值的数量。我需要处理len(users)=200000len(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

1 个答案:

答案 0 :(得分:2)

如果你想从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]))