是否有一种更简单快捷的方法来获取包含列表或 numpy 数组中相同元素的索引的索引 dict答案

【问题标题】：Is there a simpler and faster way to get an indexes dict in which contains the indexes of the same elements in a list or a numpy array是否有一种更简单快捷的方法来获取包含列表或 numpy 数组中相同元素的索引的索引 dict
【发布时间】：2018-10-26 03:16:47
【问题描述】：

说明：

我有一个包含简单整数（正且不大）的大型数组，例如 1、2、... 等。例如：[1, 1, 2, 2, 1, 2]。我想得到一个字典，其中使用列表中的单个值作为字典的键，并使用该值的索引列表作为字典的值。

问题：

有没有更简单快捷的方式在 python 中得到预期的结果？（数组可以是列表或numpy数组）

代码：

a = [1, 1, 2, 2, 1, 2]
results = indexes_of_same_elements(a)
print(results)

预期结果：

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

【问题讨论】：

标签： python arrays numpy indexing

【解决方案1】：

您可以在这里使用矢量化方法避免迭代，尤其是np.unique + np.argsort：

idx = np.argsort(a)
el, c = np.unique(a, return_counts=True)

out = dict(zip(el, np.split(idx, c.cumsum()[:-1])))

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

性能

a = np.random.randint(1, 100, 10000)

In [183]: %%timeit
     ...: idx = np.argsort(a)
     ...: el, c = np.unique(a, return_counts=True)
     ...: dict(zip(el, np.split(idx, c.cumsum()[:-1])))
     ...:
897 µs ± 41.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [184]: %%timeit
     ...: results = {}
     ...: for i, k in enumerate(a):
     ...:     results.setdefault(k, []).append(i)
     ...:
2.61 ms ± 18.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

【讨论】：

【解决方案2】：

我们可以利用元素是“简单”（即非负且不太大？）整数这一事实。

诀窍是构造一个每行只有一个元素的稀疏矩阵，然后将其转换为按列表示。这通常比argsort 快，因为如果稀疏矩阵为 MxN 且 nnz 非零，则此变换为 O(M + N + nnz)。

from scipy import sparse

def use_sprsm():
    x = sparse.csr_matrix((a, a, np.arange(a.size+1))).tocsc()
    idx, = np.where(x.indptr[:-1] != x.indptr[1:])
    return {i: a for i, a in zip(idx, np.split(x.indices, x.indptr[idx[1:]]))}

# for comparison

def use_asort():
    idx = np.argsort(a)
    el, c = np.unique(a, return_counts=True)
    return dict(zip(el, np.split(idx, c.cumsum()[:-1])))

示例运行：

>>> a = np.random.randint(0, 100, (10_000,))
>>> 
# sanity check, note that `use_sprsm` returns sorted indices
>>> for k, v in use_asort().items():
...     assert np.array_equal(np.sort(v), use_sprsm()[k])
... 
>>> timeit(use_asort, number=1000)
0.8930604780325666
>>> timeit(use_sprsm, number=1000)
0.38419671391602606

【讨论】：

太好了，我想我应该把“更简单”改为“更快”。 :) 非常感谢

【解决方案3】：

构造dict非常简单：

In []:
results = {}
for i, k in enumerate(a):
    results.setdefault(k, []).append(i)   # str(k) if you really need the key to be a str
print(results)

Out[]:
{1: [0, 1, 4], 2: [2, 3, 5]}

您也可以使用results = collections.defaultdict(list) 然后results[k].append(i) 而不是results.setdefault(k, []).append(i)

【讨论】：

感谢您的回答。 python 或 numpy 中是否有内置方法可以完成类似的工作？