python - 如何将二维 NumPy 数组的相应值映射到一维数组

Question

这段代码是我写的:

data = np.array([[3,6], [5,9], [4, 8]])

orig_x, orig_y = np.split(data, 2, axis=1)

x = np.array([3, 4])
y = np.zeros((len(x)))

for i in range(len(x)):
    y[i] = orig_y[np.where(orig_x == x[i])[0]]

基本上，我有一个 2D NumPy 数组。我将它分成两个一维数组 orig_x 和 orig_y，一个存储 x 轴的值，另一个存储 y 轴的值。

我还有另一个一维 NumPy 数组，其中包含 orig_x 数组中存在的一些值。我想找到 x 数组中每个值的 y 轴值。我使用一个简单的循环创建了这个方法，但它非常慢，因为我将它与数千个值一起使用。

你有更好的主意吗？也许通过使用 NumPy 函数？

注意:也可以为这个问题起一个更好的标题。对不起:(

Answer 1

您可以创建一个掩码来覆盖您想要的 x 列值，然后使用该掩码从 y 列中选择值。

data = np.array([[3,6], [5,9], [4, 8]])

# the values you want to lookup on the x-axis
x = np.array([3, 4])

mask = np.isin(data[:,0], x)
data[mask,1]

输出:

array([6, 8])

这里的关键函数是使用np.isin。这基本上是将 x 或 data 广播到适当的形状并进行逐元素比较:

mask = data[:,0,None] == x
y_mask = np.logical_or.reduce(mask, axis=1)
data[y_mask, 1]

输出:

array([6, 8])