将整数数组转换为索引字典

Question

我有一个（大）整数数组，比如

materials = [0, 0, 47, 0, 2, 2, 47]  # ...

只有很少的唯一条目，我想将其转换为索引字典，即，

d = {
    0: [0, 1, 3],
    2: [4, 5],
    47: [2, 6],
    }

这样做最有效的方法是什么？ （欢迎使用 NumPy。）

Answer 1

不需要numpy，这些是标准的python结构，dict理解可以很好地解决你的问题：

materials = [0, 0, 47, 0, 2, 2, 47]

d = {v : [i for i,x in enumerate(materials) if x==v] for v in set(materials)}

print(d)

结果：

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

[i for i,x in enumerate(materials) if x==v] 查找列表中元素的所有索引（index 只查找第一个）

在我的答案的第一个版本中，我在列表本身上进行迭代，但这有点浪费，因为它会在出现很多情况时多次覆盖键，并且内部列表理解具有 n 复杂性所以整体复杂度不是很好。

在我写这个最后的评论时，有人建议迭代独特的元素，这很好，所以把输入列表变成set！

Answer 2

使用enumerate() 和dict.setdefault() 函数的替代解决方案：

materials = [0, 0, 47, 0, 2, 2, 47]
d = {}
for k,m in enumerate(materials):
    d.setdefault(m, []).append(k)

print(d)

输出：

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

Answer 3

你可能会发现collections.defaultdict在这里很有用，当第一次找到一个元素时，它会为你创建一个新列表。

from collections import defaultdict

indices = defaultdict(list)

for i, elem in enumerate(materials):
    indices[elem].append(i)

Answer 4

这里有一个 numpy 解决方案：

import numpy as np

a = np.random.randint(0, 1000, 1000000)
index = np.argsort(a, kind='mergesort')
as_  = a[index]
jumps = np.r_[0, 1 + np.where(np.diff(as_) != 0)[0]]
result = {k: v for k, v in zip(as_[jumps], np.split(index, jumps[1:]))}

基准

numpy 以不太大的n 获胜；因为它使用 O(n log n) 排序算法，所以边距很小（pp2 是一种变体，它用快速排序代替了缓慢但稳定的合并排序，但代价是必须在之后对单个索引列表进行排序，pp3 将完整排序替换为argpartition 如果唯一元素的数量与元素的数量相比，这会提高一些速度。）：

原始数组中有10个不同的整数值：

原始数组中的 100 个不同的整数值：

基准代码供参考：

import numpy as np
from collections import defaultdict
import perfplot


def pp(a):
    index = np.argsort(a, kind='mergesort')
    as_ = a[index]
    jumps = np.r_[0, 1 + np.where(np.diff(as_) != 0)[0]]
    pp_out = {k: v for k, v in zip(as_[jumps], np.split(index, jumps[1:]))}
    return pp_out


def pp2(a):
    index = np.argsort(a)
    as_ = a[index]
    jumps = np.r_[0, 1 + np.where(np.diff(as_) != 0)[0]]
    pp_out = {k: np.sort(v)
              for k, v in zip(as_[jumps], np.split(index, jumps[1:]))}
    return pp_out


def Denziloe_JFFabre(a):
    df_out = {v: [i for i, x in enumerate(a) if x == v] for v in set(a)}
    return df_out


def FCouzo(a):
    fc_out = defaultdict(list)
    for i, elem in enumerate(a):
        fc_out[elem].append(i)
    return fc_out


def KKSingh(a):
    kks_out = defaultdict(list)
    list(map(lambda x: kks_out[x[0]].append(x[1]), zip(a, range(len(a)))))
    return kks_out


def TMcDonaldJensen(a):
    mdj_out = defaultdict(list)
    for i, elem in enumerate(a):
        mdj_out[elem].append(i)
    return mdj_out


def RomanPerekhrest(a):
    rp_out = {}
    for k, m in enumerate(a):
        rp_out.setdefault(m, []).append(k)
    return rp_out


def SchloemerHist(a):
    np.histogram(a, bins=np.arange(min(a), max(a)+2))
    return


def SchloemerWhere(a):
    out = {v: np.where(v == a)[0] for v in set(a)}
    return out


perfplot.show(
        setup=lambda n: np.random.randint(0, 10, n),
        kernels=[
            pp, pp2, Denziloe_JFFabre, FCouzo, KKSingh,
            TMcDonaldJensen, RomanPerekhrest, SchloemerHist, SchloemerWhere
            ],
        n_range=[2**k for k in range(19)],
        xlabel='len(a)',
        logx=True,
        logy=True,
        )

Answer 5

理解力可以很好地做到这一点：

d = {key:[i for i, v in enumerate(materials) if v == key] for key in set(materials)}

Answer 6

我会使用defaultdict，它更有效（O(n) 时间，与 Jean 的答案相比，O(n^2)）：

from collections import defaultdict
materials = [0, 0, 47, 0, 2, 2, 47]
d = defaultdict(list)
for i, elem in enumerate(materials):
    d[elem].append(i)

d 现在等于：

defaultdict(<type 'list'>, {0: [0, 1, 3], 2: [4, 5], 47: [2, 6]})

Answer 7

另一个单线，这次是numpy.where：

out = {v: np.where(v == a)[0] for v in numpy.unique(a)}

（对于某些应用程序，布尔数组可能就足够了：

out = {v: v == a for v in numpy.unique(a)}

)

请注意，对于大型数组，numpy.unique 比 set() 快，如果只有几个唯一条目，则速度要快很多。

无论如何，对于大多数数组大小，上述方法是目前最快的方法：

10 个不同的整数：

100 个不同的整数：

代码：

import numpy as np
from collections import defaultdict
import perfplot


def pp(a):
    index = np.argsort(a, kind='mergesort')
    as_ = a[index]
    jumps = np.r_[0, 1 + np.where(np.diff(as_) != 0)[0]]
    pp_out = {k: v for k, v in zip(as_[jumps], np.split(index, jumps[1:]))}
    return pp_out


def pp2(a):
    index = np.argsort(a)
    as_ = a[index]
    jumps = np.r_[0, 1 + np.where(np.diff(as_) != 0)[0]]
    pp_out = {k: np.sort(v)
              for k, v in zip(as_[jumps], np.split(index, jumps[1:]))}
    return pp_out


def Denziloe_JFFabre(a):
    df_out = {v: [i for i, x in enumerate(a) if x == v] for v in np.unique(a)}
    return df_out


def FCouzo(a):
    fc_out = defaultdict(list)
    for i, elem in enumerate(a):
        fc_out[elem].append(i)
    return fc_out


def KKSingh(a):
    kks_out = defaultdict(list)
    list(map(lambda x: kks_out[x[0]].append(x[1]), zip(a, range(len(a)))))
    return kks_out


def TMcDonaldJensen(a):
    mdj_out = defaultdict(list)
    for i, elem in enumerate(a):
        mdj_out[elem].append(i)
    return mdj_out


def RomanPerekhrest(a):
    rp_out = {}
    for k, m in enumerate(a):
        rp_out.setdefault(m, []).append(k)
    return rp_out


def SchloemerHist(a):
    np.histogram(a, bins=np.arange(min(a), max(a)+2))
    return


def SchloemerWhere(a):
    out = {v: np.where(v == a)[0] for v in np.unique(a)}
    return out


def SchloemerBooleanOnly(a):
    out = {v: v == a for v in np.unique(a)}
    return out


perfplot.show(
        setup=lambda n: np.random.randint(0, 100, n),
        kernels=[
            pp, pp2, Denziloe_JFFabre, FCouzo, KKSingh,
            TMcDonaldJensen, RomanPerekhrest, SchloemerHist, SchloemerWhere,
            SchloemerBooleanOnly
            ],
        n_range=[2**k for k in range(17)],
        xlabel='len(a)',
        logx=True,
        logy=True,
        )

Answer 8

为了好玩，这里有一个使用numpy.histogram的解决方案：

np.histogram(a, bins=np.arange(min(a), max(a)+2))

我认为它可能表现良好，但 Paul 的解决方案仍然更好：