Advanced Array Manipulation with numpy.concatenate, numpy.split

When you use numpy.concatenate, it’s tempting to think it just sticks arrays side by side, but the reality is a bit more nuanced. The key lies in the axis parameter, which dictates how the arrays are joined and how their shapes interact during this process.

Consider two arrays:

import numpy as np

a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6], [7, 8]])

Both a and b have the shape (2, 2). Now, if you concatenate along axis=0, you’re stacking them vertically:

c = np.concatenate((a, b), axis=0)
print(c)
# Output:
# [[1 2]
#  [3 4]
#  [5 6]
#  [7 8]]
print(c.shape)  # (4, 2)

The number of rows doubles because you’re appending rows from b underneath a. The number of columns stays the same because concatenation is along the rows.

Flip the axis to 1:

d = np.concatenate((a, b), axis=1)
print(d)
# Output:
# [[1 2 5 6]
#  [3 4 7 8]]
print(d.shape)  # (2, 4)

Now the arrays are joined side by side, doubling the columns and keeping the number of rows intact.

What’s critical here is that the arrays you’re concatenating must have matching dimensions except for the axis along which you’re concatenating. If they don’t, numpy throws a ValueError. This means:

• If you concatenate along axis=0, all arrays must have the same number of columns.
• If you concatenate along axis=1, all arrays must have the same number of rows.

Here’s a subtle case where things break:

a = np.array([[1, 2, 3]])
b = np.array([[4, 5]])
np.concatenate((a, b), axis=0)
# ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)

The second array b is actually 1D because it’s just one row without explicit nesting, so numpy rejects it. Fixing this requires ensuring dimensionality matches:

b = np.array([[4, 5]])
np.concatenate((a, b), axis=0)
# Works fine now because both have shape (1, n)

Another powerful aspect of numpy.concatenate comes from chaining multiple arrays and concatenating along different axes to reshape complex data structures. For example, if you have a list of arrays from a batch process:

batches = [np.random.rand(2, 3) for _ in range(5)]
full_data = np.concatenate(batches, axis=0)
print(full_data.shape)  # (10, 3)

Here, each batch has 2 rows and 3 columns, and concatenating along axis=0 stacks these batches vertically, expanding the number of samples.

But beware – numpy.concatenate itself does not reshape arrays internally. If you want to reshape the result, you need to do it explicitly:

x = np.array([[1,2],[3,4]])
y = np.array([[5,6],[7,8]])

z = np.concatenate((x, y), axis=0)
print(z.shape)  # (4, 2)

# Now reshape to a single row with 8 elements
z_reshaped = z.reshape(1, 8)
print(z_reshaped)

Understanding how the dimensions line up before and after concatenation can save you from many runtime errors. The takeaway is that concatenate is about gluing arrays along an existing axis, not altering their internal shape or dimensionality.

So, whenever you see concatenate in code, check the shapes of the arrays involved, and visualize how the axis parameter affects the final structure. This mental model of “stacking” or “lining up” along a specific dimension is fundamental.

One last thing: if you want to concatenate along a new axis, you should look at numpy.stack instead, which creates a higher-dimensional array by adding a new axis, rather than merging along an existing one. For example:

x = np.array([1, 2, 3])
y = np.array([4, 5, 6])

stacked = np.stack((x, y), axis=0)
print(stacked)
# Output:
# [[1 2 3]
#  [4 5 6]]
print(stacked.shape)  # (2, 3)

Notice how the shape changes compared to concatenate:

concatenated = np.concatenate((x, y), axis=0)
print(concatenated)
# Output: [1 2 3 4 5 6]
print(concatenated.shape)  # (6,)

This subtle difference between concatenate and stack is at the heart of many array manipulation patterns, and mastering it lets you build complex pipelines with clarity and efficiency. The key to wielding numpy.concatenate well is always to keep your eye on shapes and axes, treating arrays like Lego blocks that can only snap together if their edges match perfectly.

Once you get this down, slicing, dicing, and joining data becomes less of a chore and more of a creative process. And that’s really what working with numpy is about – bending multi-dimensional data to your will without breaking a sweat. Next up, we’ll see how to efficiently break arrays apart using numpy.split, which complements this stacking behavior by letting you divide data cleanly.

Breaking down arrays efficiently with numpy.split

numpy.split is the natural inverse of concatenation, letting you slice an array into multiple sub-arrays along a specified axis. Unlike fancy slicing with indices, it’s designed to divide arrays into equal or nearly equal chunks, which can be a huge time saver when processing data in parallel or in batches.

The function signature looks like this:

numpy.split(ary, indices_or_sections, axis=0)

Here, ary is the input array, indices_or_sections controls how the split happens, and axis is where you chop it up. The magic is in indices_or_sections: it can be either an integer or a list of indices.

If you pass an integer N, numpy.split tries to split the array into N equal parts. If the array size along the chosen axis isn’t divisible by N, it raises an error. For example:

import numpy as np

x = np.arange(12)
print(x)
# Output: [ 0  1  2  3  4  5  6  7  8  9 10 11]

# Split into 3 equal parts
parts = np.split(x, 3)
for part in parts:
    print(part)
# Output:
# [0 1 2 3]
# [4 5 6 7]
# [8 9 10 11]

If you try to split into a number that doesn’t divide the length evenly:

np.split(x, 5)
# ValueError: array split does not result in an equal division

To handle uneven splits, you use a list of indices indicating where to cut. These indices specify the boundaries between sub-arrays:

parts = np.split(x, [3, 7, 10])
for part in parts:
    print(part)
# Output:
# [0 1 2]
# [3 4 5 6]
# [7 8 9]
# [10 11]

This cuts the array at positions 3, 7, and 10, producing four sub-arrays. Notice how the last sub-array covers the remainder after the last index.

The axis argument works as you’d expect, slicing along that dimension. Here’s a 2D example:

a = np.arange(16).reshape(4,4)
print(a)
# Output:
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]
#  [12 13 14 15]]

# Split rows into 2 parts
row_parts = np.split(a, 2, axis=0)
for part in row_parts:
    print(part)
# Output:
# [[0 1 2 3]
#  [4 5 6 7]]
# [[ 8  9 10 11]
#  [12 13 14 15]]

And splitting columns:

col_parts = np.split(a, [1, 3], axis=1)
for part in col_parts:
    print(part)
# Output:
# [[ 0]
#  [ 4]
#  [ 8]
#  [12]]
# [[ 1  2]
#  [ 5  6]
#  [ 9 10]
#  [13 14]]
# [[ 3]
#  [ 7]
#  [11]
#  [15]]

Under the hood, numpy.split is just a convenient wrapper for numpy.array_split or slicing with np.take, but its strictness on equal splits makes it a good guardrail when you want predictable chunks.

When you combine numpy.split with concatenate, you get a powerful toolkit for data manipulation. For example, you might split a large dataset into batches, process them independently, then concatenate the results back:

data = np.arange(20).reshape(10, 2)

# Split into 5 batches (each with 2 rows)
batches = np.split(data, 5, axis=0)

processed = []
for batch in batches:
    # Double the values in each batch
    processed.append(batch * 2)

result = np.concatenate(processed, axis=0)
print(result)
# Output:
# [[ 0  2]
#  [ 4  6]
#  [ 8 10]
#  [12 14]
#  [16 18]
#  [20 22]
#  [24 26]
#  [28 30]
#  [32 34]
#  [36 38]]

In practice, split is invaluable for dividing data into chunks for parallel processing, cross-validation folds, or simply organizing data flow. Its strictness on equal division means you often pair it with array_split, which relaxes that constraint by allowing uneven splits.

Here’s a quick comparison to highlight that difference:

# Using array_split to split into uneven parts without error
uneven_parts = np.array_split(x, 5)
for part in uneven_parts:
    print(part)
# Output:
# [0 1 2]
# [3 4]
# [5 6]
# [7 8]
# [9 10 11]

Knowing when to use split versus array_split depends on whether your application can handle uneven chunks or needs strict uniformity.