Search code examples
numpy

How to get the NumPy array I want


How can I change the array like this?

arr = [
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],
[30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41],
...,
[130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141],
]

=> [
[ [0, 1], [20, 21], [30, 31], ,.., [130, 131]],
[ [2, 3], [22, 23], [32, 33], .., [132, 133]],
[ [4, 5], [24, 25], [34, 35], .., [134, 135]],
[ [6, 7], [26, 27], [36, 37], .., [136, 137]],
[ [8, 9], [28, 29], [38, 39], .., [138, 139]],
[ [10, 11], [30, 31], [40, 41], .., [140, 141]],
...
]

My attempt was best modified to the following, which was not the result I wanted.

[ [0, 1], [2, 3], [4, 5], ,.., [10, 11]],
[ [20, 21], [22, 23], [24, 25], .., [30, 31]],

Solution

  • You can use reshape+swapaxes:

    n, m = arr.shape
    out = arr.reshape(n, m//2, 2).swapaxes(0, 1)
    
    # or
    out = arr.reshape(n, -1, 2).swapaxes(0, 1)
    

    Example:

    # input
    arr = np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
                    [20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],
                    [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41],
                    [130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141]])
    
    # output
    array([[[  0,   1], [ 20,  21], [ 30,  31], [130, 131]],
           [[  2,   3], [ 22,  23], [ 32,  33], [132, 133]],
           [[  4,   5], [ 24,  25], [ 34,  35], [134, 135]],
           [[  6,   7], [ 26,  27], [ 36,  37], [136, 137]],
           [[  8,   9], [ 28,  29], [ 38,  39], [138, 139]],
           [[ 10,  11], [ 30,  31], [ 40,  41], [140, 141]]])