使用列表推导代替两个for循环

时间:2018-09-26 01:54:10

标签: python list-comprehension

我正在尝试一次实现:

W_jk = 0 if |j-k| > X

我尝试了这个,但是似乎容易出错。

W[i,j] = [0 if abs(i-j) > X for i in range(0,len(W)) for j in range(0,len(W[0]))]

还有其他选择吗?谢谢。

[编辑1]:W = np.random.randn(m,d)持续m,d

[编辑2]:

对于一个子问题,当j1-k1 = j2-k2且| j1-k1 | <= X时,W_j1k1 = W_j2k2。列表理解可以做到吗?

2 个答案:

答案 0 :(得分:3)

编辑:如果您原始的W数据结构是NumPy数组,则可以进行一些奇特的索引编制,以避免完全循环。

import numpy as np
i = np.arange(W.shape[0])
j = np.arange(W.shape[1])
ii, jj = np.meshgrid(j, i)  # note the flipped indices!
absval = np.abs(ii - jj)
mask = absval > X
W[mask] = 0

这比嵌套循环快得多,效率更高。

原始帖子

几项预备赛:

首先,我注意到abs(i - j) <= X没有条件,所以我加入了else 1条件。

第二,我无法确定W是否已经存在(即已预先分配),或者您是否正在理解中创建它。我把它当作存在。

第三,如果上一个问题的答案是len尚不存在,那么在理解本身中使用W是有问题的。

n_rows = 10 # change to suit your needs
n_cols = 5  # change to suit your needs
W = [ [0 if abs(i - j) > X else 1 for i in range(n_cols) ] for j in range(n_rows) ]

答案 1 :(得分:1)

您显然想将0分配给列表W的现有列表中的选择性项目。在这种情况下,使用嵌套的for循环有选择地更新项目会更有效:

for i in range(len(W)):
    for j in range(len(W[i])):
        if abs(i - j) > X:
            W[i][j] = 0

但是,如果您不关心效率而只想使其成为一行代码,则可以使用以下列表推导,以便在不满足条件|j-k| > X时重用当前值:

W = [[0 if abs(i - j) > X else W[i][j] for j in range(len(W[i]))] for i in range(len(W))]

请注意,您问题中的W[i,j]不是访问列表列表中项目的正确语法。您应该改用W[i][j]

编辑:对于您的子问题,您可以使用以下列表理解,尽管并不能通过任何方式对其进行优化:

W = [[next((W[j2][k2] for j2 in range(len(W)) for k2 in range(len(W[j2])) if j1 - k1 == j2 - k2 and abs(j1 - k1) <= X), W[j1][k1]) for k1 in range(len(W[j1]))] for j1 in range(len(W))]