NumPy Array Manipulation

Module 3: NumPy

NumPy: Array Manipulation

Numerical data does not always arrive in the shape a calculation needs. Think back to the midterm we graded in the last sessions:

  • Some professors grade exercise by exercise, so their marks arrive with one row per exercise instead of one row per student.
  • The marks of the Segovia campus arrive in a separate array, and we want a single table for the whole course.
  • The exam has three parts (open problems, debugging, multiple choice), and sometimes we want to analyse each part on its own.

In this lesson, we will learn to reorganize the same values, combine compatible arrays, and divide arrays into useful pieces. The challenge at the end puts all these tools together on the midterm.

import numpy as np

# The Madrid marks and the value of each exercise, from the last session
weights = np.array([20, 20, 10, 10, 10, 5, 5, 5, 5, 5, 5])

marks = np.array([
    [0.75, 1, 0.5, 1, 0.5, 1, 1, 0, 1, 1, 1],
    [0.5, 0.25, 1, 0.5, 0, 1, 0, 1, 1, 0, 1],
    [1, 1, 1, 1, 0.5, 1, 1, 1, 0, 1, 1],
    [0.25, 0.5, 0.5, 0.5, 0.5, 0, 1, 0, 1, 0, 1],
    [1, 1, 0.5, 1, 1, 1, 0, 1, 1, 1, 1],
    [0.5, 0.5, 0.5, 0, 1, 1, 1, 0, 0, 1, 0],
])

Reshape Arrays

Any NumPy array has access to the reshape() method. This method creates a new NumPy array with the specified shape.

# We create an array of length 6
x = np.arange(0, 6, 1)
print(x)
[0 1 2 3 4 5]
# Reshape array to 2 x 3
reshaped = x.reshape(2, 3)
print(reshaped)
print(reshaped.shape)
[[0 1 2]
 [3 4 5]]
(2, 3)

The method reshape() creates a new array, it does not modify the original!

print(f"Original:\n{x}")
print(f"Reshaped:\n{reshaped}")
Original:
[0 1 2 3 4 5]
Reshaped:
[[0 1 2]
 [3 4 5]]
Exercise: Reshape a Range

Create the following matrix without using the function np.array.

\(\begin{pmatrix} 1 & 2\\ 3 & 4\\ 5 & 6 \end{pmatrix}\)

arr = np.arange(1, 7).reshape(3, 2)
print(arr)
[[1 2]
 [3 4]
 [5 6]]
Exercise: Fill by Columns

Create the following matrix without using the function np.array.

\(\begin{pmatrix} 1 & 4\\ 2 & 5\\ 3 & 6 \end{pmatrix}\)

arr = np.arange(1, 7).reshape(2, 3).T
print(arr)
[[1 4]
 [2 5]
 [3 6]]
Exercise: Reshape and Slice

Create the following matrix without using the function np.array.

\(\begin{pmatrix} 0 & 1 & 2 & 3\\ 4 & 5 & 6 & 7\\ 8 & 9 & 10 & 11\\ 12 & 13 & 14 & 15\\ \end{pmatrix}\)

And then, using that matrix, extract:

\(\begin{pmatrix} 6 & 7\\ 10 & 11\\ \end{pmatrix}\)

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

When using reshape, you can add a single -1 to any dimension. When you do, NumPy decides the axis size to match the total number of elements in the array. For instance:

  • given an array with \(N\) elements, calling reshape(a, -1) will return a shape (a, b) where b = N / a
  • given an array with \(N\) elements, calling reshape(a, b, -1) will return a shape (a, b, c) where c = N / (a * b)

As the dimensions of reshape must always match exactly the number of elements in an array, using -1 will only work if the division produces an integer.

Reshape Is Not Transpose

In Session 14 we saw that a form filled exercise by exercise must be reshaped with one row per exercise. Here are the marks of three students in four exercises, typed exercise by exercise:

form_export = np.array([
    1, 0.5, 0.25,    # exercise 1: student 1, student 2, student 3
    0.5, 1, 0,       # exercise 2
    1, 1, 0.5,       # exercise 3
    0, 1, 1,         # exercise 4
])

by_exercise = form_export.reshape(4, 3)
print(by_exercise)
[[1.   0.5  0.25]
 [0.5  1.   0.  ]
 [1.   1.   0.5 ]
 [0.   1.   1.  ]]

We want one row per student. A tempting mistake is to reshape to (3, 4):

print(form_export.reshape(3, 4))
[[1.   0.5  0.25 0.5 ]
 [1.   0.   1.   1.  ]
 [0.5  0.   1.   1.  ]]

The shape is right, but the values are wrong: reshape always fills row by row, so the first row mixes exercise 1 and exercise 2. What we need is to swap rows and columns, the transpose we saw in Session 15:

by_student = by_exercise.T
print(by_student)
[[1.   0.5  1.   0.  ]
 [0.5  1.   1.   1.  ]
 [0.25 0.   0.5  1.  ]]
Exercise: Spot the Difference

Without running any code, what is the first row of form_export.reshape(3, 4) and what is the first row of by_exercise.T? Which one contains the exam of student 1?

[1.   0.5  0.25 0.5 ]
[1.  0.5 1.  0. ]

The first row of the reshape is the first four values of the form: the three marks of exercise 1 and the first mark of exercise 2. The first row of the transpose is the first column of by_exercise: exercise 1 to 4 of student 1. Only the transpose contains the exam of student 1.


Flatten Arrays

Another useful method is flatten(). It will create a one-dimensional array from a multi-dimensional array, effectively collapsing all the dimensions into a single dimension.

# We create a matrix 3x3
x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(f"Original:\n{x}")
Original:
[[1 2 3]
 [4 5 6]
 [7 8 9]]
# Flatten array
flattened = x.flatten()
print(f"Flattened:\n{flattened}")
Flattened:
[1 2 3 4 5 6 7 8 9]

Joining Arrays

In data processing and analysis, you often need to combine arrays. This can be done in several ways depending on the structure and the desired outcome. Below, we’ll explore some common methods for joining arrays using NumPy.

Concatenation

Concatenation is the process of joining two or more arrays along an existing axis.

Concatenation on axis 0:

\(\text{concat}_0(\mathbf{a}, \mathbf{b}) = [a_0, a_1, b_0, b_1]\)

\(\text{concat}_0(A, B) = \begin{pmatrix} a_{00} & a_{01} \\ a_{10} & a_{11} \\ b_{00} & b_{01} \\ b_{10} & b_{11} \\ \end{pmatrix}\)

# Create two 2-D arrays
array1 = np.array([[1, 2], [3, 4]])
array2 = np.array([[5, 6], [7, 8]])
print("Original shapes", array1.shape, "&", array2.shape)

# Concatenate along rows (axis=0)
concatenated_rows = np.concatenate((array1, array2), axis=0)
print("\nConcatenated along rows:\n", concatenated_rows)
print("\nConcatenated shape:", concatenated_rows.shape)
Original shapes (2, 2) & (2, 2)

Concatenated along rows:
 [[1 2]
 [3 4]
 [5 6]
 [7 8]]

Concatenated shape: (4, 2)

When you use np.concatenate(a, b, axis=0), it combines two arrays by stacking them vertically, one on top of the other. Imagine you have two bricks, and you place one brick directly on top of the other.

Now, let’s compare this to np.sum(a, axis=0). When you sum an array along axis=0, you are adding up the elements in each column. It’s like pressing down on the columns and “squeezing” the values together into a single row.

array1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print("Original array:\n", array1)

result = np.sum(array1, axis=0)
print("We 'squeeze' the values down:\n  |  |  |\n  V  V  V\n", result)
Original array:
 [[1 2 3]
 [4 5 6]
 [7 8 9]]
We 'squeeze' the values down:
  |  |  |
  V  V  V
 [12 15 18]

Concatenation on axis 1:

\(\text{concat}_1(\mathbf{a}, \mathbf{b}) = \text{ERROR}\)

\(\text{concat}_1(A, B) = \begin{pmatrix} a_{00} & a_{01} & b_{00} & b_{01} \\ a_{10} & a_{11} & b_{10} & b_{11} \\ \end{pmatrix}\)

# Create two 2-D arrays
array1 = np.array([[1, 2], [3, 4]])
array2 = np.array([[5, 6], [7, 8]])
print("Original shapes", array1.shape, "+", array2.shape)

# Concatenate along columns (axis=1)
concatenated_columns = np.concatenate((array1, array2), axis=1)
print("\nConcatenated along columns:\n", concatenated_columns)
print("\nConcatenated shape:", concatenated_columns.shape)
Original shapes (2, 2) + (2, 2)

Concatenated along columns:
 [[1 2 5 6]
 [3 4 7 8]]

Concatenated shape: (2, 4)

When you use np.concatenate(a, b, axis=1), it combines two arrays by stacking them horizontally, side to side. Imagine you have two bricks, and you place one brick at the left of the other.

Compare this to np.sum(a, axis=1). When you sum an array along axis=1, you are adding up the elements in each row. It’s like “squeezing” the values together into a single column.

array1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print("Original array:\n", array1)

result = np.sum(array1, axis=1)
print("We 'squeeze' the values sideways:\n  ->\n  -> \n  ->\n", result)
Original array:
 [[1 2 3]
 [4 5 6]
 [7 8 9]]
We 'squeeze' the values sideways:
  ->
  -> 
  ->
 [ 6 15 24]

Pay attention to the following errors:

# You cannot concatenate on a non-existing axis
array1 = np.array([[1, 2], [3, 4]])
array2 = np.array([[5, 6], [7, 8]])
arr = np.concatenate((array1, array2), axis=2)
# The arrays must have the same number of dimensions
arr1 = np.array([[1, 2, 3], [4, 5, 6]])
arr2 = np.array([7, 8, 9])
arr = np.concatenate((arr1, arr2), axis=0)
print(arr)

Stacking

Stacking is another way to join arrays, but it can add an extra dimension to the result.

Horizontal stack: Stack arrays in sequence horizontally (column wise).

\(\text{hstack}(\mathbf{a}, \mathbf{b}) = [a_0, a_1, b_0, b_1]\)

\(\text{hstack}(A, B) = \begin{pmatrix} a_{00} & a_{01} & b_{00} & b_{01} \\ a_{10} & a_{11} & b_{10} & b_{11} \\ \end{pmatrix}\)

This is equivalent to concatenation along the second axis, except for 1-D arrays where it concatenates along the first axis.

# Create two 1-D arrays
array1 = np.array([1, 2, 3, 4])
array2 = np.array([5, 6, 7, 8])
print("Original shapes", array1.shape, "+", array2.shape)

# Horizontal stacking
hstacked = np.hstack((array1, array2))
print("\nHorizontally stacked:\n", hstacked)
print("\nHorizontally stacked shape:", hstacked.shape)
Original shapes (4,) + (4,)

Horizontally stacked:
 [1 2 3 4 5 6 7 8]

Horizontally stacked shape: (8,)
# Create two 2-D arrays
array1 = np.array([[1, 2], [3, 4]])
array2 = np.array([[5, 6], [7, 8]])
print("Original shapes", array1.shape, "+", array2.shape)

# Horizontal stacking
hstacked = np.hstack((array1, array2))
print("\nHorizontally stacked:\n", hstacked)
print("\nHorizontally stacked shape:", hstacked.shape)
Original shapes (2, 2) + (2, 2)

Horizontally stacked:
 [[1 2 5 6]
 [3 4 7 8]]

Horizontally stacked shape: (2, 4)

Vertical stack: Stack arrays in sequence vertically (row wise).

\(\text{vstack}(\mathbf{a}, \mathbf{b}) = \begin{pmatrix} a_0 & a_1 \\ b_0 & b_1 \\ \end{pmatrix}\)

\(\text{vstack}(A, B) = \begin{pmatrix} a_{00} & a_{01} \\ a_{10} & a_{11} \\ b_{00} & b_{01} \\ b_{10} & b_{11} \\ \end{pmatrix}\)

This is equivalent to concatenation along the first axis after 1-D arrays of shape (N,) have been reshaped to (1,N).

# Create two 1-D arrays
array1 = np.array([1, 2, 3, 4])
array2 = np.array([5, 6, 7, 8])
print("Original shapes", array1.shape, "+", array2.shape)

# Vertical stacking
vstacked = np.vstack((array1, array2))
print("\nVertically stacked:\n", vstacked)
print("\nVertically stacked shape:", vstacked.shape)
Original shapes (4,) + (4,)

Vertically stacked:
 [[1 2 3 4]
 [5 6 7 8]]

Vertically stacked shape: (2, 4)
# Create two 2-D arrays
array1 = np.array([[1, 2], [3, 4]])
array2 = np.array([[5, 6], [7, 8]])
print("Original shapes", array1.shape, "+", array2.shape)

# Horizontal stacking
vstacked = np.vstack((array1, array2))
print("\nVertically stacked:\n", vstacked)
print("\nVertically stacked shape:", vstacked.shape)
Original shapes (2, 2) + (2, 2)

Vertically stacked:
 [[1 2]
 [3 4]
 [5 6]
 [7 8]]

Vertically stacked shape: (4, 2)
# vstack is equilavent to ...

# Create two 1-D arrays
array1 = np.array([1, 2, 3, 4])
array2 = np.array([5, 6, 7, 8])
print("Original shapes", array1.shape, "+", array2.shape)

# Reshape the arrays to (1, 4)
array1 = array1.reshape(1, -1)
array2 = array2.reshape(1, -1)
print("Reshaped", array1.shape, "+", array2.shape)

# Concatenate along axis 1
concatenated_columns = np.concatenate((array1, array2), axis=0)
print("\nConcatenated along columns:\n", concatenated_columns)
print("\nConcatenated shape:", concatenated_columns.shape)
Original shapes (4,) + (4,)
Reshaped (1, 4) + (1, 4)

Concatenated along columns:
 [[1 2 3 4]
 [5 6 7 8]]

Concatenated shape: (2, 4)
Exercise: Rouché-Frobenius Theorem

A non-homogeneous system of linear equations \(Ax=b\) with \(n\) variables has a solution if and only if the rank of its coefficient matrix \(A\) is equal to the rank of its augmented matrix \([A|b]\). If there are solutions, we can find that:

  • if \(rank(A) = n\), the solution is unique,
  • if \(rank(A) < n\), there are infinite solutions.

Write a function rouche_frobenius_theorem(A, b) that builds the augmented matrix \([A|b]\) with a stacking function and returns a message saying whether the system is inconsistent, has a unique solution, or has infinitely many solutions. Use np.linalg.matrix_rank() to compute ranks.

Test it with:

A = np.array([[2, 1, -1], [1, 3, 2], [1, -1, 2]])
b = np.array([8, 13, 3])
def rouche_frobenius_theorem(A, b):
    # Turn vector b with shape (N,) into a column with shape (N, 1).
    B = b.reshape(-1, 1)
    Ab = np.hstack((A, B))
    rank_A = np.linalg.matrix_rank(A)
    rank_Ab = np.linalg.matrix_rank(Ab)

    if rank_A == rank_Ab:
        if rank_A == A.shape[1]:
            return "The system is consistent and has a unique solution."
        return "The system is consistent and has infinitely many solutions."
    else:
        return "The system is inconsistent."

# Example usage
A = np.array([[2, 1, -1], [1, 3, 2], [1, -1, 2]])
b = np.array([8, 13, 3])

result = rouche_frobenius_theorem(A, b)
print(result)
The system is consistent and has a unique solution.

The vector b must be reshaped into a column before np.hstack() can place it next to A.

Stack

Column stacking is used to stack 1-D arrays as columns into a 2-D array.

\(\text{column_stack}(\mathbf{a}, \mathbf{b}) = \begin{pmatrix} a_0 & b_0 \\ a_1 & b_1 \\ \end{pmatrix}\)

Take a sequence of 1-D arrays and stack them as columns to make a single 2-D array. 2-D arrays are stacked as-is, just like with hstack. 1-D arrays are turned into 2-D columns first.

# Create two 1-D arrays
array1 = np.array([1, 2, 3])
array2 = np.array([4, 5, 6])

# Column stack
column_stacked = np.column_stack((array1, array2))
print("Column stacked:\n", column_stacked)
Column stacked:
 [[1 4]
 [2 5]
 [3 6]]

Row stacking is used to stack 1-D arrays as rows into a 2-D array.

\(\text{row_stack}(\mathbf{a}, \mathbf{b}) = \begin{pmatrix} a_0 & a_1 \\ b_0 & b_1 \\ \end{pmatrix}\)

This is equivalent to concatenation along the first axis after 1-D arrays of shape (N,) have been reshaped to (1,N). In older code you will find it as np.row_stack(), an alias of np.vstack() that is deprecated since NumPy 2.0. Use np.vstack() instead.

# Create two 1-D arrays
array1 = np.array([1, 2, 3])
array2 = np.array([4, 5, 6])

# Row stack
row_stacked = np.vstack((array1, array2))
print("Row stacked:\n", row_stacked)
Row stacked:
 [[1 2 3]
 [4 5 6]]

Back to the midterm: a student took the exam late, and we want to add her marks as a new row.

late_student = np.array([1, 0.75, 1, 0.5, 1, 1, 0, 1, 1, 1, 1])

with_late = np.vstack((marks, late_student))
print(with_late.shape)
(7, 11)

Splitting Arrays

Splitting arrays is the process of dividing an array into multiple sub-arrays. This can be useful for various data manipulation tasks. Below, we’ll explore some common methods for splitting arrays using NumPy.

Split

The np.split() function splits an array into multiple sub-arrays along a specified axis.

# Create an array
arr = np.array([1, 2, 3, 4, 5, 6])

# Split the array into 3 equal parts
split_array = np.split(arr, 3)

for idx in range(len(split_array)):
    print(f"Split {idx + 1}:\n{split_array[idx]}")
Split 1:
[1 2]
Split 2:
[3 4]
Split 3:
[5 6]

Array Split

The np.array_split() function allows you to split an array into unequal parts if needed.

For an array of length l that should be split into n sections, it returns l % n sub-arrays of size l//n + 1 and the rest of size l//n.

# Split the array into 4 parts (unequal)
split_array_unequal = np.array_split(arr, 4)

for idx in range(len(split_array_unequal)):
    print(f"Split {idx + 1}:\n{split_array_unequal[idx]}")
Split 1:
[1 2]
Split 2:
[3 4]
Split 3:
[5]
Split 4:
[6]

Horizontal Split

The np.hsplit() function splits an array horizontally (column-wise).

# Create a 2D array
array_2d = np.array([[1, 2, 3], [4, 5, 6]])

# Split the 2D array into 3 columns
hsplit_array = np.hsplit(array_2d, 3)

for idx in range(len(hsplit_array)):
    print(f"Split {idx + 1}:\n{hsplit_array[idx]}")
Split 1:
[[1]
 [4]]
Split 2:
[[2]
 [5]]
Split 3:
[[3]
 [6]]

Vertical Split

The np.vsplit() function splits an array vertically (row-wise).

# Split the 2D array into 2 rows
vsplit_array = np.vsplit(array_2d, 2)

for idx in range(len(vsplit_array)):
    print(f"Split {idx + 1}:\n{vsplit_array[idx]}")
Split 1:
[[1 2 3]]
Split 2:
[[4 5 6]]

Splitting at Given Positions

Instead of a number of equal parts, all the split functions also accept a list of positions where to cut. The three parts of the midterm have different sizes: exercises 1 and 2, exercises 3 to 5, and exercises 6 to 11. We cut the columns before index 2 and before index 5:

open_part, debugging_part, mc_part = np.hsplit(marks, [2, 5])

print(open_part.shape, debugging_part.shape, mc_part.shape)
(6, 2) (6, 3) (6, 6)
Exercise: Split the Weights

Split weights into the same three parts, and check that the parts are worth 40, 30, and 30 points.

open_w, debugging_w, mc_w = np.split(weights, [2, 5])

print(np.sum(open_w), np.sum(debugging_w), np.sum(mc_w))
40 30 30

weights is one-dimensional, so we use np.split() (or np.hsplit(), which behaves the same for 1D arrays).


Challenge: One Table for the Whole Course

Deactivate AI assistant tools, and try this challenge on your own.

The professor in Segovia graded the midterm exercise by exercise, and sent you the marks of their 4 students as one long array: the 4 marks of exercise 1, then the 4 marks of exercise 2, and so on.

segovia_form = np.array([
    0.5, 1, 0, 0.75,
    1, 1, 0.25, 0.5,
    0.5, 1, 0.5, 1,
    1, 0.5, 0.5, 1,
    1, 1, 0, 0.5,
    1, 1, 0, 1,
    1, 0, 1, 1,
    1, 1, 0, 0,
    1, 1, 1, 0,
    0, 1, 0, 1,
    1, 1, 0, 1,
])

Without writing any loop:

  1. Turn segovia_form into a table with one row per student and one column per exercise.
  2. Join it with the Madrid marks into a single table, course, with the Madrid students first. Print its shape.
  3. Split course into the three parts of the exam, and compute the points of each student in each part (use weights).
  4. Compute the final grade of every student, out of 10.
  5. Build a report with two columns: the student number (1 to 10) and the final grade.
# 1. One row per exercise first, then transpose
segovia = segovia_form.reshape(11, 4).T
print(segovia)

# 2. Join along the rows
course = np.concatenate((marks, segovia), axis=0)  # or np.vstack()
print(course.shape)

# 3. Split the columns at the start of each part
open_part, debugging_part, mc_part = np.hsplit(course, [2, 5])
open_w, debugging_w, mc_w = np.split(weights, [2, 5])

print(np.sum(open_part * open_w, axis=1))
print(np.sum(debugging_part * debugging_w, axis=1))
print(np.sum(mc_part * mc_w, axis=1))

# 4. Final grades
final_grades = course @ weights / 10
print(final_grades)

# 5. Report
report = np.column_stack((np.arange(1, 11), final_grades))
print(report)
[[0.5  1.   0.5  1.   1.   1.   1.   1.   1.   0.   1.  ]
 [1.   1.   1.   0.5  1.   1.   0.   1.   1.   1.   1.  ]
 [0.   0.25 0.5  0.5  0.   0.   1.   0.   1.   0.   0.  ]
 [0.75 0.5  1.   1.   0.5  1.   1.   0.   0.   1.   1.  ]]
(10, 11)
[35. 15. 40. 15. 40. 20. 30. 40.  5. 25.]
[20. 15. 25. 15. 25. 15. 25. 25. 10. 25.]
[25. 20. 25. 15. 25. 15. 25. 25. 10. 20.]
[8.  5.  9.  4.5 9.  5.  8.  9.  2.5 7. ]
[[ 1.   8. ]
 [ 2.   5. ]
 [ 3.   9. ]
 [ 4.   4.5]
 [ 5.   9. ]
 [ 6.   5. ]
 [ 7.   8. ]
 [ 8.   9. ]
 [ 9.   2.5]
 [10.   7. ]]

Note that segovia_form.reshape(4, 11) would also give the right shape, but with the marks mixed up.


Summary

You do not need to memorize every joining and splitting function. Focus on the shape of each input, the shape you want as output, and the axis along which the operation acts.

  • reshape() changes the arrangement of values but not their count.
  • flatten() returns a one-dimensional copy.
  • concatenate() joins arrays along an existing axis.
  • stack(), vstack(), and hstack() provide common arrangements.
  • split() requires equal pieces; array_split() also permits unequal pieces.

Before running any joining operation, predict the output shape. Then explain why all non-joining dimensions must match.


Homework and Review

Work through the homework exercises available here.

Homework