使用字典元组键的最后两个元素求和

Question

我有名单

A = [(i,j,k,l,m)]
B = [(l,m,k)]

和字典

C = {(i,j,k,l,m): val}
D = {(l,m,k): other_val}

我想创建一个E字典，以便

E = {(i,j,k): C[(i,j,k,l,m)]*D[(l,m,k)]}

假设列表和字典中的所有索引约定都匹配。我有下面的非肾盂，非常缓慢的解决方案。对于非常大的A大小（例如，500万行），是否有任何Pythonic方法可以快速执行此操作

E = {}
for i,j,k,l,m in A:
    E[i,j,k] = sum(
        C[i,j,k,l,m] * D[l2,m2,k2] 
        for l2,m2,k2 in B if l2==l and m2==m and k2==k)

下面是生成样本数据集的代码，该样本数据集接近要处理的实际大小

import numpy as np
np.random.seed(1)

Irange = range(50)
Jrange = range(10)
Krange = range(80)
Lrange = range(8)
Mrange = range(18)

A = [
    (i,j,k,l,m)
    for i in Irange
    for j in Jrange
    for k in Krange
    for l in Lrange
    for m in Mrange]
B = [
    (l,m,k)
    for k in Krange
    for l in Lrange
    for m in Mrange]

C = {key: np.random.uniform(1,10) for key in A}

D = {key: np.random.uniform(0,1) for key in B}

E = {}
for i,j,k,l,m in A:
    E[i,j,k] = sum(
        C[i,j,k,l,m] * D[l2,m2,k2]
        for l2,m2,k2 in B if l2==l and m2==m and k2==k)

Answer 1

from collections import defaultdict

b = set(B)  # O(#B)
E = defaultdict(float)
for i,j,k,l,m in A: # O(#A)
    if (l, m, k) in b:
        E[i,j,k] += C[i,j,k,l,m] * D[l, m, k]

这种方法具有O(#A + #B)复杂性。 撇开正确性问题不谈，天真的实现是O(#A * #B)

Answer 2

我正在发布我的足够快的解决方案。如果你仍然看到改进的可能性，我很乐意测试一下。（我希望一些库已经有了一个更快的解决方案；也许有，但我的问题/解决方法不够清晰，无法使用）

以下是数据生成代码：

import numpy as np
from datetime import datetime
np.random.seed(1)

Irange = range(50)
Jrange = range(10)
Krange = range(80)
Lrange = range(8)
Mrange = range(18)

A = [
    (i,j,k,l,m)
    for i in Irange
    for j in Jrange
    for k in Krange
    for l in Lrange
    for m in Mrange]
B = [
    (l,m,k)
    for k in Krange
    for l in Lrange
    for m in Mrange]

C = {key: np.random.uniform(1,10) for key in A}

D = {key: np.random.uniform(0,1) for key in B}

首先，启动计时器并引入一个列表unique_ijk：

start_timer = datetime.now() #Start counting time
unique_ijk = list(set([(i,j,k) for i,j,k,l,m in A]))

然后，创建一个名为lm_given_ijk的字典，该字典使用与给定的i，j，k元组键对应的l，m索引列表进行赋值

lm_given_ijk = {(i,j,k):[] for i,j,k in unique_ijk}
for i,j,k,l,m in A:
    lm_given_ijk[i,j,k].append((l,m))

最后，按如下方式使用lm_given_ijk来创建E

E = {(i,j,k): sum(C[i,j,k,l,m]*D[l,m,k] for l,m in lm_given_ijk[i,j,k]) 
                  for i,j,k in unique_ijk}
print("Elapsed time is %s seconds.\n"%(datetime.now()-start_timer).total_seconds())

输出：

Elapsed time is 6.446798 seconds.

写下所有这些，我同意评论说这是一个numpy数组的事情。它可以提高速度，但我对6.4秒感到满意

Answer 3

D与B具有相同的键，那么为什么要对B进行迭代并对一个元素求和？您可以直接获取D[l,m,k]

E[i,j,k] = C[i,j,k,l,m] * D[l,m,k]

转换为听写理解：

E = {(i,j,k): C[i,j,k,l,m]*D[l,m,k] for i,j,k,l,m in A}

在我的电脑上，这给了大约150倍的速度（15秒）→ 0.086s），但出于实用性考虑，我将您的所有输入维度减半，因为您的原始代码运行了一分钟多，没有生成任何输出