将三个for循环写入一行(python)

时间:2016-12-16 18:34:08

标签: python nested-loops

我有以下程序

import numpy as np

X = np.asarray([ [2,3,4,5,1,1,1],
  [1,2,3,4,1,1,1],
  [3,3,3,3,1,1,1]])

Nrows=X.shape[0]
Ncols1=X.shape[1]

cols=np.asarray( [[1,3]])
Ncols2=cols.shape[1]

Y = np.ndarray(shape=(Nrows,Ncols2))

for i in range(Nrows):
    for j in range(Ncols2):
        sum=0
        for k in range(Ncols1):
            sum=sum+np.exp(-(X[i,k]-X[i,cols[0,j])**2)
        Y[i,j]=sum

我想将三个循环缩减为一行。这里的动机部分是为了加速考虑,部分是为了练习python技能。

所以我可以将它写成两个循环:

for i in range(Nrows):
    for j in range(Ncols1):
        Y[i,j]=np.sum(np.exp(-(X[i,]-X[i,cols[0,j]])**2))

我可以把它写成一个循环:

for i in range(Nrows):
   Y[i,]=map(lambda j: np.sum(np.exp(-(X[i,]-X[i,cols[0,j]])**2)),range(Ncols2))

但我无法弄清楚如何将其变成一行。我的直觉是尝试

Y[:,]=map(lambda j: np.sum(np.exp(-(X[:,]-X[:,cols[0,j]])**2)),range(Ncols2))

但是失败了" ValueError:操作数无法与形状一起广播(3,7)(3,)"

非常感谢任何帮助。

-Nick

1 个答案:

答案 0 :(得分:1)

Y=np.asarray([np.sum(np.exp(-(X[i,]-X[i, cols[0,j]])**2)) for i in range(Nrows) for j in range(Ncols2)]).reshape((Nrows,Ncols2))

>>> Y
array([[ 1.80902144,  1.38631883],
       [ 2.85771284,  1.38668872],
       [ 4.05494692,  4.05494692]])

这使用list comp,它几​​乎只是你在一行上的for循环。

[np.sum(np.exp(-(X[i,]-X[i, cols[0,j]])**2)) for i in range(Nrows) for j in range(Ncols2)]

这有点与此相符:

Y=[]
for i in range(Nrows):
  for j in range(Ncols2):
    Y.append(#your big formula that uses i and j as arguments)

因为它是一个列表,你必须从中创建一个数组,并根据自己的喜好对其进行重新整形,因此列表comp之前为np.asarray,而

之后为.reshape

永远不会玩numpy太多...我喜欢它!