python - 如何从两个相关的 NumPy 数组填充字典 - 替代追加

Question

我正在尝试学习如何使用 NP 数组而不是列表。但是，当最初创建 NP 数组时不知道属于 NP 数组的所有值时，我正在努力确定将值添加到 numpy 数组的正确方法。

下面是我之前代码的简化示例，其中我根据值所属的类别(在本例中为国家/地区)将值附加到列表中。因此值列表中的第 3 个值属于国家列表中的第 3 个国家(230 属于法国)。

country = ['Germany','USA','France','Germany','USA','USA']
values = [250,220,230,180,180,220]

cv = {'Germany':[],'USA':[],'France':[]}

for i in range(len(country)):
    cv[country[i]].append(values[i])

这是一个简化的示例，实际上我需要附加更长的值，因此我想知道重写代码以得到 numpy 数组而不是列表的最有效方法是什么.

Answer 1

处理大型数组的最快方法是使用 boolean masking .

首先，将初始列表包装到 NumPy 数组中:

countries = ['Germany','USA','France','Germany','USA','USA']
values = [250,220,230,180,180,220]
countries = np.array(countries)
values = np.array(values)

现在对于生成的字典 cv 中的每个唯一国家立即填充值而不附加:

cv = {country: values[countries == country]
      for country in set(countries)}
print(cv)

将给予:

{'France': array([230]),
 'USA': array([220, 180, 220]),
 'Germany': array([250, 180])}

---

时间:

为了一个更具代表性的例子，我增加了数组的大小:

n = 500
countries = np.array(['Germany','USA','France','Germany','USA','USA'] * n)
values = np.array([250,220,230,180,180,220] * n)

# My version
%%timeit
cv = {country: values[countries == country]
      for country in set(countries)}
# 987 µs ± 5.46 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# MathiasRa's version
%%timeit
cv = {'Germany':[],'USA':[],'France':[]}
for i in range(len(countries)):
    cv[countries[i]].append(values[i])
# 1.67 ms ± 9.94 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# jpp's version
%%timeit
d = defaultdict(list)
for c, v in zip(countries, values):
    d[c].append(v)
res = {k: np.array(v) for k, v in d.items()}
# 1.5 ms ± 2.37 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Johann Bzh's version
%%timeit
d={}
for c, v in zip(countries, values):
    tmp_vec = np.array([v])
    try:
        d[c] = np.hstack((d[c],tmp_vec))
    except KeyError: #first occurence of this country
        d[c] = tmp_vec  
# 13.3 ms ± 125 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)