python - 从整数列表创建索引字典

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:])}

sample 运行-

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 如何帮助减少元素？

使用基于排序的方法 #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 索引，这仍然使它相对更高效。