python - "Checkbord-like"二进制选择

Question

我有一个二维数组

test = np.arange(25).reshape((5, 5))
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

为此我需要一个像选择器一样的“棋盘”。基本上，我想选择第一行中的奇数元素，第二行中的偶数元素，第三行中的奇数元素等。

在这种情况下(因为行号是奇数)我可以通过

test.flatten()[::2]
Out[22]: array([ 0,  2,  4,  6,  8, 10, 12, 14, 16, 18, 20, 22, 24])

这正是两个维度中的“所有其他元素”。但如果我们尝试这样做

test2 = np.arange(16).reshape((4, 4))
test2
Out[23]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])
test2.flatten()[::2]
Out[27]: array([ 0,  2,  4,  6,  8, 10, 12, 14])

不出所料，它没有成功。

无论(2d)矩阵的形状如何，生成我想要的选择类型的通用方法是什么？是否可以轻松(甚至可能)扩展到 3d 矩阵？

我确实更喜欢高效(矢量化或基于 cython)的方法。

<小时/>

奖励:3D:

test3 = np.arange(25*3).reshape((3, 5, 5))
Out[30]: 
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24]],
       [[25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49]],
       [[50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59],
        [60, 61, 62, 63, 64],
        [65, 66, 67, 68, 69],
        [70, 71, 72, 73, 74]]])

这里不是关于奇数和偶数行，而是关于相邻元素。 两个相邻元素不应具有相同的颜色(属于同一 bool 选择器)。

也就是说，如果我们选择 test3[0, ...].flatten()[::2] 我们会得到 [0, ... 24] 。第一维中 0 的相邻元素是 25，所以我们不希望这样。预期输出是

np.hstack((test3[0, ...].flatten()[::2], test3[1, ...].flatten()[1::2], test3[2, ...].flatten()[::2], ))
Out[42]: 
array([ 0,  2,  4,  6,  8, 10, 12, 14, 17, 19, 21, 23, 25, 27, 29, 31, 32,
       34, 36, 38, 40, 42, 44, 46])

Answer 1

这是一种矢量化解决方案 -

def checker_select(a):
    m,n = a.shape[-2:]
    i,j = np.ogrid[:m,:n]
    mask = (i+j)%2==0
    return a.reshape(-1,m,n)[:,mask].ravel()

使用数组初始化可能会更快 -

def checker_select_v2(a):
    m,n = a.shape[-2:]
    mask = np.zeros((m,n), dtype=bool)
    mask[::2,::2] = 1
    mask[1::2,1::2] = 1
    return a.reshape(-1,m,n)[:,mask].ravel()

2D 的示例运行 -

In [117]: a
Out[117]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

In [118]: checker_select(a)
Out[118]: array([ 0,  2,  4,  6,  8, 10, 12, 14, 16, 18, 20, 22, 24])

3D 的示例运行 -

In [144]: a
Out[144]: 
array([[[  0,   1,   2,   3,   4],
        [  5,   6,   7,   8,   9],
        [ 10,  11,  12,  13,  14],
        [ 15,  16,  17,  18,  19],
        [ 20,  21,  22,  23,  24]],

       [[ 25,  26,  27,  28,  29],
        [ 30,  31,  32,  33,  34],
        [ 35,  36,  37,  38,  39],
        [ 40,  41,  42,  43,  44],
        [ 45,  46,  47,  48,  49]],

        ....

In [145]: checker_select(a)
Out[145]: 
array([  0,   2,   4,   6,   8,  10,  12,  14,  16,  18,  20,  22,  24,
        25,  27,  29,  31,  33,  35,  37,  39,  41,  43,  45,  47,  49,
        ....