在 ndarrays 中查找最接近零的值

Question

我有两个名为 A 和 B 的 numpy ndarray。每个 ndarray 的维度为 2 x 3。对于每个网格点，我必须找到两个数组中最接近零的元素并相应地分配一个标志。该标志对数组 A 取值 1，对数组 B 取值 2。也就是说，如果数组 A 的 (0,0) 中的元素（即第 0 行和第 0 列）与 (0,0) 相比最接近于零数组 B 的元素，则输出在第 0 行和第 0 列的位置分配一个值 1。输出数组的维度为 1 x 3。

下面我举个例子

A= np.array([[0.1,2,0.3],[0.4,3,2]])
B= np.array([[1,0.2,0.5],[4,0.03,0.02]])

输出应该是

[[1,2,1],[1,2,2]]

有没有一种不用写 for 循环的有效方法？非常感谢。

Answer 1

我会这样做：

import numpy as np

a = np.array([[0.1,2,0.3],[0.4,3,2]])
b = np.array([[1,0.2,0.5],[4,0.03,0.02]])

c = np.abs(np.stack([a, b])).argmin(0)+1

输出：

array([[1, 2, 1],
       [1, 2, 2]])