Python：使用list-comprehension优化For循环

Question

我最近才开始使用Python的列表推导。我对他们并不满意。如何使用列表推导优化下面的代码？

heatmap_color = []

for r in xrange(len(heatmap)):
    heatmap_color.append([])
    for c in xrange(len(heatmap[r])):
        heatmap_color[r].append(cmap.to_rgba(heatmap[r][c], alpha=0.5, bytes=True))

此代码的作用是创建一个2D数组（表示彩色图像的像素值）。最后一行将热图中的每个现有灰度像素映射到alpha通道= 0.5的颜色，以便可以在heatmap_color数组的相应条目中输入这些颜色。

目前运行时间不到30秒。我想尽可能减少这段时间。

谢谢！

Answer 1

嵌套循环：

heatmap_color = [[cmap.to_rgba(r, alpha=0.5, bytes=True) for r in c] for c in heatmap]

请注意，我们不需要使用索引（由xrange()生成）;而是直接在

上循环。

你也可以在这里使用map()和functools.partial()来加快速度：

from functools import partial
heatmap_rgba = partial(cmap.to_rgba, alpha=0.5, bytes=True)
heatmap_color = [map(heatmap_rgba, c) for c in heatmap]

甚至：

from functools import partial
heatmap_rgba = partial(cmap.to_rgba, alpha=0.5, bytes=True)
heatmap_per_list = partial(map, heatmap_rgba)
heatmap_color = map(heatmap_per_list, heatmap)

因为map()可能比列表推导更快，尤其是与functools.partial一起使用时，以避免lambda需要额外的堆叠推送。