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

該代碼的作用是創建一個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要快，尤其是與functools.partial一起使用時，可以避免lambda所需的額外堆棧推。