从整数列表创建索引字典

Question

我有一个(长)数组a的一些不同的整数.我现在想创建一个字典,其中键是整数,值是索引数组,其中a出现相应的整数.这个

import numpy

a = numpy.array([1, 1, 5, 5, 1])
u = numpy.unique(a)
d = {val: numpy.where(a==val)[0] for val in u}

print(d)

{1: array([0, 1, 4]), 5: array([2, 3])}

工作正常,但第一次打电话似乎相当浪费unique,接下来是几个where.

np.digitize 似乎不是理想的,因为你必须提前指定垃圾箱.

有关如何改进上述的任何想法？

Answer 1

方法#1

基于排序的一种方法是 -

def group_into_dict(a): 
    # Get argsort indices
    sidx = a.argsort()

    # Use argsort indices to sort input array
    sorted_a = a[sidx]

    # Get indices that define the grouping boundaries based on identical elems
    cut_idx = np.flatnonzero(np.r_[True,sorted_a[1:] != sorted_a[:-1],True])

    # Form the final dict with slicing the argsort indices for values and
    # the starts as the keys
    return {sorted_a[i]:sidx[i:j] for i,j in zip(cut_idx[:-1], cut_idx[1:])}

样品运行 -

In [55]: a
Out[55]: array([1, 1, 5, 5, 1])

In [56]: group_into_dict(a)
Out[56]: {1: array([0, 1, 4]), 5: array([2, 3])}

阵列上的时间与1000000元素和不同比例的唯一数字比较建议的与原始的一个 -

# 1/100 unique numbers
In [75]: a = np.random.randint(0,10000,(1000000))

In [76]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
1 loop, best of 3: 6.62 s per loop

In [77]: %timeit group_into_dict(a)
10 loops, best of 3: 121 ms per loop

# 1/1000 unique numbers
In [78]: a = np.random.randint(0,1000,(1000000))

In [79]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
1 loop, best of 3: 720 ms per loop

In [80]: %timeit group_into_dict(a)
10 loops, best of 3: 92.1 ms per loop

# 1/10000 unique numbers
In [81]: a = np.random.randint(0,100,(1000000))

In [82]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
10 loops, best of 3: 120 ms per loop

In [83]: %timeit group_into_dict(a)
10 loops, best of 3: 75 ms per loop

# 1/50000 unique numbers
In [84]: a = np.random.randint(0,20,(1000000))

In [85]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
10 loops, best of 3: 60.8 ms per loop

In [86]: %timeit group_into_dict(a)
10 loops, best of 3: 60.3 ms per loop

因此,如果您正在处理的是20或多或少的唯一数字,请坚持阅读原始数据 ; 否则基于排序似乎运作良好.

方法#2

Pandas 基于一个适合很少的唯一数字 -

In [142]: a
Out[142]: array([1, 1, 5, 5, 1])

In [143]: import pandas as pd

In [144]: {u:np.flatnonzero(a==u) for u in pd.Series(a).unique()}
Out[144]: {1: array([0, 1, 4]), 5: array([2, 3])}

阵列上的计时与1000000具有20独特元素的元素 -

In [146]: a = np.random.randint(0,20,(1000000))

In [147]: %timeit {u:np.flatnonzero(a==u) for u in pd.Series(a).unique()}
10 loops, best of 3: 35.6 ms per loop

# Original solution
In [148]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
10 loops, best of 3: 58 ms per loop

并减少独特元素 -

In [149]: a = np.random.randint(0,10,(1000000))

In [150]: %timeit {u:np.flatnonzero(a==u) for u in pd.Series(a).unique()}
10 loops, best of 3: 25.3 ms per loop

In [151]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
10 loops, best of 3: 44.9 ms per loop

In [152]: a = np.random.randint(0,5,(1000000))

In [153]: %timeit {u:np.flatnonzero(a==u) for u in pd.Series(a).unique()}
100 loops, best of 3: 17.9 ms per loop

In [154]: %timeit {val: np.where(a==val)[0] for val in np.unique(a)}
10 loops, best of 3: 34.4 ms per loop

如何pandas减少元素的帮助？

基于排序approach #1,对于20独特元素的情况,获取argsort索引是瓶颈 -

In [164]: a = np.random.randint(0,20,(1000000))

In [165]: %timeit a.argsort()
10 loops, best of 3: 51 ms per loop

现在,pandas基于函数为我们提供了唯一的元素,无论是负数还是任何东西,我们只是简单地与输入数组中的元素进行比较,而无需进行排序.让我们看看这方面的改进:

In [166]: %timeit pd.Series(a).unique()
100 loops, best of 3: 3.17 ms per loop

当然,它需要获得np.flatnonzero指数,这仍然使其相对更有效率.