具有嵌套 for 循环、条件和累加器的列表理解

Question

我正在尝试将这段代码转换为列表理解：

a = np.random.rand(10) #input vector
n = len(a) # element count of input vector
b = np.random.rand(3) #coefficient vector
nb = len(b) #element count of coefficients
d = nb #decimation factor (could be any integer < len(a))
 
c = []
for i in range(0, n, d):
    psum = 0
    for j in range(nb):
        if i + j < n:
            psum += a[i + j]*b[j]
    c.append(psum)

我尝试了以下建议：

• 带有累加器的列表理解
• 嵌套 for 循环以列出具有不同“if”条件的理解

例如：

from itertools import accumulate
c = [accumulate([a[i + j] * b[j] for j in range(nb) if i + j < n] ) for i in range(0, n, d)]

后来，当试图从c （例如c[:index] ）获取值时：

TypeError: 'NoneType' object is not subscriptable

或者：

from functools import partial
def get_val(a, b, i, j, n):
    if i + j < n:
        return(a[i + j] * b[j])
    else:
        return(0)
c = [
         list(map(partial(get_val, i=i, j=j, n=n), a, b)) 
             for i in range(0, n, d) 
             for j in range(nb)
    ]

在get_val中，返回(a[i + j] * b[j])

IndexError: invalid index to scalar variable.

或者：

psum_pieces = [[a[i + j] * b[j] if i + j < n else 0 for j in range(nb)] for i in range(0, n, d)]
c = [sum(psum) for psum in psum_pieces]

以及这些方法的许多其他迭代。 任何指导将不胜感激。

Answer 1

如果我理解正确你想要的是

res = [sum(a[i+j]*b[j] for j in range(nb) if i+j < n) for i in range(0,n,d)]

对于每个i ，这将在结果列表中添加 j 的乘积总和a[i+j]*b[j]当i+j < n j从0变化到nb-1

Answer 2

你真的不需要在这里使用列表理解。 使用 numpy，您可以创建一个不直接在解释器中运行任何循环的快速流水线解决方案。

首先将a转换为形状为(n // d, nb)的二维数组。 缺少的元素（即循环中的i + j >= n ）可以为零，因为这将使相应的增量psum ：

# pre-compute i+j as a 2D array
indices = np.arange(nb) + np.arange(0, n, d)[:, None]
# we only want valid locations
mask = indices < n

t = np.zeros(indices.shape)
t[mask] = a[indices[mask]]

现在您可以直接计算c为

(t * b).sum(axis=1)

我怀疑如果你将此解决方案与任何用 vanilla python 编写的没有用 numba 编译的东西进行基准测试，它会快得多。