时间戳

首页 > 解决方案 > Numpy 数组以索引为值的字典

python - Numpy 数组以索引为值的字典

问题描述

我有一个带有整数值的 numpy 数组，我们称这个数组为 x。

我想创建某种列表，其中对于每个值，我都有 x 的索引来保存这个值。

例如，对于：

x = [1,2,2,4,7,1,1,7,16]

我想得到：

{1: [0,5,6], 2:[1,2], 4:[3], 7:[4,7], 15:[16]}

我使用的括号是任意的，我不在乎我使用哪种数据结构，只要我能尽快将结果输出到文件即可。最后，我想要一个 .txt 文件，内容如下：

0,5,6

1,2

3

4,7

16

标签： pythonnumpy

解决方案

由于您提到您对值的数据结构并不挑剔，因此要获得类似于您在问题中发布的字典的内容，您可以对值中的唯一值进行字典x理解np.where：

>>> {i:np.where(x == i)[0] for i in set(x)}

{1: array([0, 5, 6]),
 2: array([1, 2]),
 4: array([3]),
 7: array([4, 7]),
 16: array([8])}

将此与通过列表进行更普通的循环进行比较，对于较大的数组，这将明显更快：

def list_method(x):
    res = {i:[] for i in set(x)}
    for i, value in enumerate(x):
         res[value].append(i)
    return res

def np_method(x):
    return {i:np.where(x == i)[0] for i in set(x)}

x = np.random.randint(1, 50, 1000000)


In [5]: %timeit list_method(x)
259 ms ± 4.03 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [6]: %timeit np_method(x)
120 ms ± 4.15 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

推荐阅读