python：使數組索引生成更有效/更優雅

Question

我正在嘗試使用python獲取一些數組索引。 此刻，代碼看起來非常麻煩，我想知道我是效率低下還是使用非Python風格。 因此，我有一個n維數組，並且嘗試生成一些索引，如下所示。 這是一個隔離的代碼示例，演示了我正在嘗試做的事情。 我將在一個簡單的2D代碼段中進行演示，但該數組可以是任意維度。

import numpy as np

a = np.random.rand(5, 5)
shape = a.shape

for i in range(len(shape)):
    shape_temp = np.zeros(len(shape), dtype=np.int)
    shape_temp[i] = 1
    p = np.meshgrid(*[np.arange (shape_temp[l], shape[l]) for l in range(len(shape))])

    # Do something with p

我想知道是否有一種更優雅，更有效的方法來生成這些索引？ 特別是，使用shape_temp變量看起來很丑，尤其是每次循環時如何創建它。

Answer 1

shape_temp = np.zeros_like(shape)

您可以使用以下表達式避免shape_temp ：

[np.arange(1 if l==i else 0, e) for l,e in enumerate(shape)]

更漂亮或更有效還是值得商bat的

另一個片段

temp = np.eye(len(shape))
[np.arange(j,k) for j,k in zip(temp[i,:],shape)]

meshgrid的替代meshgrid是mgrid或ogrid （盡管我必須更改indexing才能獲得匹配）。 表達式更緊湊，因為它們采用slices而不是ranges 。 在內部，這些函數使用arange 。

meshgrid(*[np.arange(j,k) for j,k in zip(temp[i,:],shape)],indexing='ij')
np.mgrid[[slice(j,k) for j,k in zip(temp[i,:],shape)]]

meshgrid(*[np.arange(j,k) for j,k in zip(temp[i,:],shape)],sparse=True,indexing='ij')
np.ogrid[[slice(j,k) for j,k in zip(temp[i,:],shape)]]