python - 添加 (row, col) 作为元素级 numpy 函数的参数？

Question

Numpy 非常自然地进行元素操作。然而，在编写函数时，我经常希望获得元素的位置，例如:

def neighbor_count(arr, row, col):
 hood = arr[row-1:row+1, col-1, col+1]
 return np.sum(hood)-arr[row,col]

我可以创建一个嵌套的 for 循环:

for (i, row) in enumerate(arr[1:-1]):
 for (j, col) in enumerate (row[1:-1]):
  c = neighbor_count(arr, i+1, j+1)

但这感觉就像放弃并回到我在 C++ 中的工作方式。

有什么聪明的 Pythonic 解决方案吗？

Answer 1

你想要numpy.ndenumerate ，它的工作方式与 enumerate 类似，但对于 ndarray 结构。

我无法为您的特定用例提供准确的答案，因为尚不清楚您如何尝试组合结果值，但基本思想是这样的。

for (r, c), i in np.ndenumerate(arr):
  do_something()

请注意，ndenumerate 返回一个迭代器，因此您习惯的所有 itertools 功能都应该适用于它。