以更好、更短的方式重写一个双层循环。

Question

请问以下代码有没有更好的写法？我想要计算在一个 (x, y) 网格中的函数值 z = f(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])

- nos

Answer 1

是的。你在问题中呈现的代码很好。

永远不要认为几行代码很“好看”或“酷”。重要的是清晰、易读和可维护性。其他人应该能够理解你的代码（当你需要找到一个 bug 时，你也应该在 12 个月内理解它）。

许多程序员，特别是年轻的程序员，认为“聪明”的解决方案是可取的。其实不是这样的。这就是 Python 社区的优秀之处。我们比其他社区更少地受到这种错误的影响。

- Lennart Regebro

---

我点赞了这个帖子，因为这是我自己正在学习的一课，但如果它们只是列表，那么我认为我的解决方案比 OP 的代码更好，因为它避开了不必要的迭代器。在相同的假设下，wheatie 的答案更好。 - aaronasterling

迭代器仍然存在。它并没有绕过任何东西。话虽如此，我当然经常使用列表推导式，但如果我必须嵌套它们，我会将它们展开，这样它们就更容易阅读。 - Lennart Regebro

Answer 2

你可以这样做：

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

- aaronasterling

Answer 3

你可以使用 itertools 的 product:

[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)]

如果您想要一个关于 xr 和 yr 的函数，您还可以预设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

- wheaties

---

值得注意的是，所有这些解决方案都会生成一个一维列表，而不是像原帖中的解决方案生成一个嵌套列表。 - aaronasterling

Answer 4

如果你想一次性设置所有内容，可以使用列表推导式；

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

如果您需要使用已经存在的 z，那么您无法这样做，您的代码也是最整洁的。补充一下：我不知道 lingrid 是用来做什么的，但如果它生成了一个包含100个元素的列表，则使用 aaronasterling 的列表推导式；如果您不需要创建额外的迭代器，则没有必要这样做。

- Chris Morgan

Answer 5

这显示了一般结果。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]]

- hughdbrown

Answer 6

我认为这是你正在寻找的一行代码

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

- Muhammed Aslam

Answer 7

你的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：

import numpy as 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) 作为输入。但是一般来说，如果没有巨大的性能差异 ¹ ，那么请使用您和其他人容易理解代码的方法。

---

¹ 如果你的 a 和 b 是 numpy.array，那么性能差异会很大，因为如果不需要进行 list<->numpy.array 转换，numpy 可以更快地处理数组。

- MSeifert