在 numpy 中从具有索引的 2D 矩阵构建 3D 布尔矩阵

Posted 2023-02-14

【中文标题】在 numpy 中从具有索引的 2D 矩阵构建 3D 布尔矩阵

【英文标题】：Building a 3D boolean matrix from a 2D matrix with indices, in numpy

【发布时间】：2022-01-22 07:41:13

【问题描述】：

我有一个形状为 (3, 4) 的二维矩阵，其索引范围从 0 到 8：

a = array([[0, 4, 1, 2],
           [5, 0, 2, 3],
           [8, 6, 0, 5]])

目前，我使用for 循环来构建一个形状为(9, 3, 4) 的3D 布尔数组，该数组将True 存储在每个索引的位置，对于0 到8 之间的每一行：

b = np.zeros((9, 3, 4), dtype=bool)
for i in range(9):
    b[i] = np.where(a == i, True, False)

有没有办法在没有迭代的情况下达到相同的结果，也许使用 numpy 函数？

【参考方案1】：

这是你要找的东西吗？

import numpy as np

a = np.array([[0, 4, 1, 2],
              [5, 0, 2, 3],
              [8, 6, 0, 5]])
y, x = np.mgrid[0:a.shape[0], 0:a.shape[1]]

data = np.zeros((9,) + a.shape, dtype=bool)
data[a, y, x] = True

【参考方案2】：

利用 numpy 广播的非常简短的解决方案：

b = np.array([a]*9) == np.arange(9).reshape(-1,1,1)

输出：

>>> b
array([[[ True, False, False, False],
        [False,  True, False, False],
        [False, False,  True, False]],

       [[False, False,  True, False],
        [False, False, False, False],
        [False, False, False, False]],

       [[False, False, False,  True],
        [False, False,  True, False],
        [False, False, False, False]],

       [[False, False, False, False],
        [False, False, False,  True],
        [False, False, False, False]],

       [[False,  True, False, False],
        [False, False, False, False],
        [False, False, False, False]],

       [[False, False, False, False],
        [ True, False, False, False],
        [False, False, False,  True]],

       [[False, False, False, False],
        [False, False, False, False],
        [False,  True, False, False]],

       [[False, False, False, False],
        [False, False, False, False],
        [False, False, False, False]],

       [[False, False, False, False],
        [False, False, False, False],
        [False, False, False, False]]])