np.expand_dims() and np.squeeze() in NumPy - Time & Space Complexity

Time Complexity: np.expand_dims() and np.squeeze()

O(1)

Understanding Time Complexity

We want to understand how the time to run np.expand_dims() and np.squeeze() changes as the size of the input array grows.

Specifically, how does adding or removing dimensions affect the work done?

Scenario Under Consideration

Analyze the time complexity of the following code snippet.

import numpy as np

arr = np.random.rand(1000, 10)
expanded = np.expand_dims(arr, axis=1)
squeezed = np.squeeze(expanded, axis=1)

This code creates a 2D array, adds a new dimension, then removes that dimension.

Identify Repeating Operations

Identify the loops, recursion, array traversals that repeat.

• Primary operation: Creating views of the array shape without copying data.
• How many times: The operations happen once each, no loops over elements.

How Execution Grows With Input

Adding or removing dimensions does not involve touching each element, so the time stays almost the same as input grows.

Input Size (n)	Approx. Operations
10	Very few operations, just shape change
100	Still very few operations, no element-wise work
1000	Same as before, no increase in element processing

Pattern observation: The time does not grow with input size because no element-wise processing happens.

Final Time Complexity

Time Complexity: O(1)

This means the time to add or remove a dimension stays constant no matter how big the array is.

Common Mistake

[X] Wrong: "Adding or removing dimensions must take longer for bigger arrays because the data changes."

[OK] Correct: These functions only change the shape metadata, not the actual data, so the time does not depend on array size.

Interview Connect

Knowing that some numpy operations are just shape changes helps you explain efficient data handling in real projects.

Self-Check

"What if we used a function that copies the array data instead of just changing shape? How would the time complexity change?"