使嵌套'for'循環更加pythonic

Question

我是python的新手，我想知道如何通過避免顯式嵌套'for'循環並使用python的隱式循環來提高效率。 我正在使用圖像數據，在這種情況下嘗試加速我的k-means算法。 這是我正在嘗試做的一個示例：

# shape of image will be something like 140, 150, 3
num_sets, rows_per_set, num_columns = image_values.shape

for set in range(0, num_sets):
    for row in range(0, rows_per_set):
        pos = np.argmin(calc_euclidean(rgb_[set][row], means_list)
        buckets[pos].append(image_values[set][row])

我今天所擁有的東西效果很好，但我想讓它更有效率。

非常感謝您的反饋和建議。

Answer 1

這是一個矢量化解決方案。 我幾乎可以肯定我的尺寸混亂（3不是真正的列數，不是嗎？），但原則應該是可識別的：

為了演示，我只將（平面）索引收集到桶中的集合和行中。

import numpy as np

k = 6
rgb_=np.random.randint(0, 9, (140, 150, 3))
means_list = np.random.randint(0, 9, (k, 3))

# compute distance table; use some algebra to leverage highly optimised
# dot product
squared_dists = np.add.outer((rgb_*rgb_).sum(axis=-1),
                             (means_list*means_list).sum(axis=-1)) \
    - 2*np.dot(rgb_, means_list.T)
# find best cluster
best = np.argmin(squared_dists, axis=-1)

# find group sizes
counts = np.bincount(best.ravel())
# translate to block boundaries
bnds = np.cumsum(counts[:-1])
# group indices by best cluster; argpartition should be
# a bit cheaper than argsort
chunks = np.argpartition(best.ravel(), bnds)
# split into buckets
buckets = np.split(chunks, bnds)

# check

num_sets, rows_per_set, num_columns = rgb_.shape

def calc_euclidean(a, b):
    return ((a-b)**2).sum(axis=-1)

for set in range(0, num_sets):
    for row in range(0, rows_per_set):
        pos = np.argmin(calc_euclidean(rgb_[set][row], means_list))
        assert pos == best[set, row]
        assert rows_per_set*set+row in buckets[pos]