匹配一个向量和一个矩阵的元素，返回一个新的矩阵

Question

我有以下 numpy arrays：

[
  [96.6074, 73.3443, 37.2574],
  [11.2753, 39.8242, 42.1898],
  [46.5815, 76.2492, 71.7398],
  [4.99222, 7.64707, 6.5171],
]

[
  True,
  False,
  False,
  True
]

现在我希望 boolean 数组的元素对应于浮点数组的行。 我想“匹配”它们并创建 2 个新的 arrays，其中：

• arr1将仅包含 floatArray 的行，其中 boolArray 中的相同行（按索引）具有True
• arr2将仅包含 floatArray 的行，其中 boolArray 中的相同行（按索引）具有False

我认为可以通过这样的循环来完成：

arr1 = np.array([])
arr2 = np.array([])

for idx, element in  np.ndenumerate(boolArray):
    if (element == True):
        # append floatArray[idx] to arr1
    else:
        # append floatArray[idx] to arr2

但我是 python 的新手，我不确定该怎么做。有什么想法吗？

Answer 1

您可以简单地使用数组掩码：

>>> arr = np.array([
        [96.6074, 73.3443, 37.2574],
        [11.2753, 39.8242, 42.1898],
        [46.5815, 76.2492, 71.7398],
        [4.99222, 7.64707, 6.5171],
    ])
>>> mask = np.array([True, False, False, True])
>>> arr[mask]
array([[96.6074 , 73.3443 , 37.2574 ],
       [ 4.99222,  7.64707,  6.5171 ]])
>>> arr[~mask]
array([[11.2753, 39.8242, 42.1898],
       [46.5815, 76.2492, 71.7398]])