python - Python：使用列表理解优化 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))

这段代码所做的是创建一个二维数组（表示彩色图像的像素值）。最后一行将 heatmap 中的每个现有灰度像素映射到 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()andfunctools.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()可以比列表推导更快，尤其是在使用 withfunctools.partial以避免额外的堆栈推送 a时lambda。