以更好,更短的方式重写双循环

时间:2010-12-17 23:38:49

标签: python loops

我想知道以下代码是否可以以更好的方式编写。基本上,我想为z = f(x, y) meshgrid计算(x, y)

a = linspace(0, xr, 100)                                                                  
b = linspace(0, yr, 100)                                                                  

for i in xrange(100):
   for j in xrange(100):
      z[i][j] = f(a[i],b[j])

7 个答案:

答案 0 :(得分:22)

呀。您在问题中提供的代码很不错。

不要以为少数几行“好”或“酷”。重要的是清晰度,可读性和可维护性。其他人应该能够理解你的代码(你应该在12个月内理解它,当你需要找到一个bug时)。

许多程序员,特别是年轻程序员,认为“聪明”的解决方案是可取的。他们不是。这就是python社区的优点。我们比其他人更少受到这种错误的折磨。

答案 1 :(得分:6)

你可以做点什么

z = [[f(item_a, item_b) for item_b in b] for item_a in a]

答案 2 :(得分:4)

您可以使用itertools'产品:

[f(i,j) for i,j in product( a, b )]

如果你真的想将这5行缩小为1,那么:

[f(i,j) for i,j in product( linspace(0,xr,100), linspace(0,yr,100)]

如果你想要一个xryr的功能,你可以将0和100的范围预设为其他东西,这样可以更进一步:

def ranged_linspace( _start, _end, _function ):
    def output_z( xr, yr ):
        return [_function( i, j ) for i,j in product( linspace( _start, xr, _end ), linspace( _start, yr, _end ) )]
    return output_z

答案 3 :(得分:2)

如果您一次全部设置,则可以使用列表理解;

[[f(a[i], b[j]) for j in range(100)] for i in range(100)]

但是,如果您需要使用已经存在的z,那么您就无法做到这一点,而您的代码是关于您将获得的最新代码。

添加:我不知道这个lingrid做了什么,但如果它产生一个100个元素的列表,请使用aaronasterling的列表理解;如果你不需要,就没有必要创建额外的迭代器。

答案 4 :(得分:0)

这显示了一般结果。 a列为6长的列表,b为4长。结果是一个包含6个列表的列表,每个嵌套列表的长度为4个元素。

>>> def f(x,y):
...     return x+y
... 
>>> a, b = list(range(0, 12, 2)), list(range(0, 12, 3))
>>> print len(a), len(b)
6 4
>>> result = [[f(aa, bb) for bb in b] for aa in a]
>>> print result
[[0, 3, 6, 9], [2, 5, 8, 11], [4, 7, 10, 13], [6, 9, 12, 15], [8, 11, 14, 17], [10, 13, 16, 19]]

答案 5 :(得分:0)

我认为这是您要寻找的一行代码

z =  [[a+b for b in linspace(0,yr,100)] for a in linspace(0,xr,100)]

答案 6 :(得分:0)

您的linspace实际上看起来可能是np.linspace。如果它是你可以在numpy数组上操作而不必明确迭代:

z = f(x[:, np.newaxis], y)

例如:

>>> import numpy as np
>>> x = np.linspace(0, 9, 10)
>>> y = np.linspace(0, 90, 10)
>>> x[:, np.newaxis] + y  # or f(x[:, np.newaxis], y)
array([[  0.,  10.,  20.,  30.,  40.,  50.,  60.,  70.,  80.,  90.],
       [  1.,  11.,  21.,  31.,  41.,  51.,  61.,  71.,  81.,  91.],
       [  2.,  12.,  22.,  32.,  42.,  52.,  62.,  72.,  82.,  92.],
       [  3.,  13.,  23.,  33.,  43.,  53.,  63.,  73.,  83.,  93.],
       [  4.,  14.,  24.,  34.,  44.,  54.,  64.,  74.,  84.,  94.],
       [  5.,  15.,  25.,  35.,  45.,  55.,  65.,  75.,  85.,  95.],
       [  6.,  16.,  26.,  36.,  46.,  56.,  66.,  76.,  86.,  96.],
       [  7.,  17.,  27.,  37.,  47.,  57.,  67.,  77.,  87.,  97.],
       [  8.,  18.,  28.,  38.,  48.,  58.,  68.,  78.,  88.,  98.],
       [  9.,  19.,  29.,  39.,  49.,  59.,  69.,  79.,  89.,  99.]])

但您也可以使用np.ogrid代替两个linspace

将numpy导入为np

>>> x, y = np.ogrid[0:10, 0:100:10]
>>> x + y  # or f(x, y)
array([[ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 1, 11, 21, 31, 41, 51, 61, 71, 81, 91],
       [ 2, 12, 22, 32, 42, 52, 62, 72, 82, 92],
       [ 3, 13, 23, 33, 43, 53, 63, 73, 83, 93],
       [ 4, 14, 24, 34, 44, 54, 64, 74, 84, 94],
       [ 5, 15, 25, 35, 45, 55, 65, 75, 85, 95],
       [ 6, 16, 26, 36, 46, 56, 66, 76, 86, 96],
       [ 7, 17, 27, 37, 47, 57, 67, 77, 87, 97],
       [ 8, 18, 28, 38, 48, 58, 68, 78, 88, 98],
       [ 9, 19, 29, 39, 49, 59, 69, 79, 89, 99]])

这在某种程度上取决于你f是什么。如果它包含math.sin等函数,则需要将其替换为numpy.sin

如果它不是numpy,那么你应该坚持使用你的选项,或者在循环时选择使用enumerate

for idx1, ai in enumerate(a):
   for idx2, bj in enumerate(b):
      z[idx1][idx2] = f(ai, bj)

这样做的好处是您不需要对range(或xrange)进行硬编码或使用len(a)作为输入。但总的来说,如果没有巨大的性能差异 1 ,那么使用您和其他人使用您的代码的方法将很容易理解。

1 如果您的abnumpy.array,则会出现明显的性能差异,因为如果没有{numpy可以更快地处理数组' <1}}&lt; - &gt; list次转换是必需的。