2

I have a case where I have input array like this below

array([[[ 1.,  0.,  2.,  0.,  3.,  0.,  4.,  0.,  5.],
        [ 6.,  0.,  7.,  0.,  8.,  0.,  9.,  0., 10.],
        [11.,  0., 12.,  0., 13.,  0., 14.,  0., 15.]],

       [[16.,  0., 17.,  0., 18.,  0., 19.,  0., 20.],
        [21.,  0., 22.,  0., 23.,  0., 24.,  0., 25.],
        [26.,  0., 27.,  0., 28.,  0., 29.,  0., 30.]]])

and I would like to get an output like the one below.

array([[[ 1.,  0.,  3.,   0.,  5.],
        [ 6.,  0.,  8.,   0., 10.],
        [11.,  0., 13.,   0., 15.]],

       [[16.,  0., 18.,   0., 20.],
        [21.,  0., 23.,   0., 25.],
        [26., 0., 28.,    0., 30.]]])

I would love if the solution can be generic not just to this example.

Barmar
  • 741,623
  • 53
  • 500
  • 612
  • Please provide enough code so others can better understand or reproduce the problem. – Community Oct 04 '22 at 17:38
  • This is not a serious or generic solution. **This approach is reading from invalid memory** (but probably doesn't crash for the small example array). I just wanted to know if it is possible. `np.lib.stride_tricks.as_strided(arr, (2,3,3,2), (8*27, 8*9, 8*4, 8)).reshape(2,3,6)[:,:,:5].copy()`. – Michael Szczesny Oct 04 '22 at 19:24

1 Answers1

3

Since the length of the last dimension cannot be guaranteed to be even, here I choose to build the bool indices:

>>> mask = np.arange(ar.shape[-1]) // 2 % 2 == 0   # np.arange() & 2 == 0 is faster
>>> mask
array([ True,  True, False, False,  True,  True, False, False,  True])
>>> ar[:, :, mask]  # or ar[..., mask]
array([[[ 1.,  0.,  3.,  0.,  5.],
        [ 6.,  0.,  8.,  0., 10.],
        [11.,  0., 13.,  0., 15.]],

       [[16.,  0., 18.,  0., 20.],
        [21.,  0., 23.,  0., 25.],
        [26.,  0., 28.,  0., 30.]]])

If the length of the last dimension can be guaranteed to be even, reshape with slicing is another technique:

>>> ar
array([[[ 1.,  0.,  2.,  0.,  3.,  0.,  4.,  0.],
        [ 6.,  0.,  7.,  0.,  8.,  0.,  9.,  0.],
        [11.,  0., 12.,  0., 13.,  0., 14.,  0.]],

       [[16.,  0., 17.,  0., 18.,  0., 19.,  0.],
        [21.,  0., 22.,  0., 23.,  0., 24.,  0.],
        [26.,  0., 27.,  0., 28.,  0., 29.,  0.]]])
>>> shape = ar.shape[:-1]
>>> ar.reshape(*shape, -1, 2)[..., ::2, :].reshape(*shape, -1)
array([[[ 1.,  0.,  3.,  0.],
        [ 6.,  0.,  8.,  0.],
        [11.,  0., 13.,  0.]],

       [[16.,  0., 18.,  0.],
        [21.,  0., 23.,  0.],
        [26.,  0., 28.,  0.]]])
Mechanic Pig
  • 6,756
  • 3
  • 10
  • 31
  • 1
    Very nice! For the particular case of "take 2, leave 2", `np.arange(a.shape[-1]) & 2 == 0` is a bit (10%-20% in a timetest benchmark) faster :-) (and for "take 4, leave 4", you can do `& 3`, etc.) – AKX Oct 04 '22 at 17:53