python - 快速将带有索引的 numpy 数组转换为以该索引为键的 numpy 数组的字典

Question

我有一组 numpy 数组。其中之一是“键”列表，我想将数组重新排列成以该键为键的数组字典。我当前的代码是:

for key, val1, val2 in itertools.izip(keys, vals1, vals2):
    dict1[key].append(val1)
    dict2[key].append(val2)

这非常慢，因为涉及的数组有数百万个条目长，而且这种情况发生了很多次。是否有可能以矢量化形式重写它？一组可能的 key 是提前知道的，并且有大约 10 个不同的 key 。

编辑:如果有 k 个不同的键并且列表有 n 长，则当前答案为 O(nk)(对每个键迭代一次)和 O(n log n)(先排序).不过，我仍在寻找 O(n) 矢量化解决方案。这是有希望的；毕竟，最简单的非向量化事物(即我已经拥有的事物)是 O(n)。

Answer 1

执行此操作的矢量化方法可能需要您对键进行排序。基本思想是对键和值进行排序以匹配。然后，每当排序键数组中有新键时，您就可以拆分 val 数组。矢量化代码看起来像这样:

import numpy as np

keys = np.random.randint(0, 10, size=20)
vals1 = np.random.random(keys.shape)
vals2 = np.random.random(keys.shape)

order = keys.argsort()
keys_sorted = keys[order]

# Find uniq keys and key changes
diff = np.ones(keys_sorted.shape, dtype=bool)
diff[1:] = keys_sorted[1:] != keys_sorted[:-1]
key_change = diff.nonzero()[0]
uniq_keys = keys_sorted[key_change]

vals1_split = np.split(vals1[order], key_change[1:])
dict1 = dict(zip(uniq_keys, vals1_split))

vals2_split = np.split(vals2[order], key_change[1:])
dict2 = dict(zip(uniq_keys, vals2_split))

由于 argsort 步骤，此方法的复杂度为 O(n * log(n))。实际上，argsort 非常快，除非 n 非常大。在 argsort 变得明显变慢之前，您可能会使用此方法遇到内存问题。