Splitting NumPy arrays

NumPy arrays can be split horizontally, vertically, and depth-wise. First, let’s create an array:

import numpy as np

a = np.arange(36).reshape(6, 6)
a

#out
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]])

Horizontal splitting of NumPy arrays

The following code splits the array a into 3 parts (horizontally) of equal sizes:

np.hsplit(a, 3)

#out
[array([[ 0,  1],
        [ 6,  7],
        [12, 13],
        [18, 19],
        [24, 25],
        [30, 31]]), array([[ 2,  3],
        [ 8,  9],
        [14, 15],
        [20, 21],
        [26, 27],
        [32, 33]]), array([[ 4,  5],
        [10, 11],
        [16, 17],
        [22, 23],
        [28, 29],
        [34, 35]])]

You can also use the split() method specifying the axis.

np.split(a, 3, axis = 1)

#out
[array([[ 0,  1],
        [ 6,  7],
        [12, 13],
        [18, 19],
        [24, 25],
        [30, 31]]), array([[ 2,  3],
        [ 8,  9],
        [14, 15],
        [20, 21],
        [26, 27],
        [32, 33]]), array([[ 4,  5],
        [10, 11],
        [16, 17],
        [22, 23],
        [28, 29],
        [34, 35]])]

Vertical splitting of NumPy arrays

The following code splits the array a into 3 parts (vertically) of equal sizes:

np.vsplit(a, 3)

#out
[array([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11]]), array([[12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23]]), array([[24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35]])]

You can also use the split() method specifying the axis.

np.split(a, 3, axis = 0)

#out
[array([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11]]), array([[12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23]]), array([[24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35]])]

Depth-wise splitting of NumPy arrays

Depth-wise splitting requires an array of minimum rank 3. Let’s create one:

import numpy as np

a = np.arange(27).reshape(3, 3, 3)
a

#out
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]]])

Use dsplit() to split a NumPy array depth-wise.

np.dsplit(a, 3)

#out
[array([[[ 0],
         [ 3],
         [ 6]],
 
        [[ 9],
         [12],
         [15]],
 
        [[18],
         [21],
         [24]]]), array([[[ 1],
         [ 4],
         [ 7]],
 
        [[10],
         [13],
         [16]],
 
        [[19],
         [22],
         [25]]]), array([[[ 2],
         [ 5],
         [ 8]],
 
        [[11],
         [14],
         [17]],
 
        [[20],
         [23],
         [26]]])]

Leave a Reply