Sorting along axes in NumPy - Time & Space Complexity
Start learning this pattern below
Jump into concepts and practice - no test required
Sorting data is common in data science to organize information. When sorting along an axis in numpy, we want to know how the time needed grows as the data size grows.
We ask: How does sorting time change when we have bigger arrays or sort along different axes?
Analyze the time complexity of the following code snippet.
import numpy as np
arr = np.random.rand(1000, 1000)
sorted_arr = np.sort(arr, axis=1)
This code creates a 1000x1000 array of random numbers and sorts each row independently.
Identify the loops, recursion, array traversals that repeat.
- Primary operation: Sorting each row of the 2D array.
- How many times: Sorting is done once per row, so 1000 times for 1000 rows.
Sorting each row takes time depending on the row length. Doing this for all rows adds up.
| Input Size (n x m) | Approx. Operations |
|---|---|
| 10 x 10 | 10 rows x sorting 10 items each ≈ 10 x 10 log 10 |
| 100 x 100 | 100 rows x sorting 100 items each ≈ 100 x 100 log 100 |
| 1000 x 1000 | 1000 rows x sorting 1000 items each ≈ 1000 x 1000 log 1000 |
Pattern observation: The time grows roughly with the number of rows times the sorting cost per row, which depends on the row length times its logarithm.
Time Complexity: O(n * m log m)
This means sorting each of the n rows of length m takes time proportional to n times m log m.
[X] Wrong: "Sorting the whole 2D array is just O(n log n) because it's one big sort."
[OK] Correct: Sorting along an axis sorts each row separately, so the time depends on sorting each row, not the whole array as one list.
Understanding how sorting time grows with data size helps you explain performance in real tasks. It shows you can think about how algorithms work on multi-dimensional data.
"What if we sorted along axis=0 (columns) instead of axis=1? How would the time complexity change?"
Practice
axis parameter control in numpy.sort?Solution
Step 1: Understand the role of
Theaxisin sortingaxisparameter tells numpy which direction to sort: 0 means sort each column, 1 means sort each row.Step 2: Differentiate from other parameters
Sorting algorithm type and order are controlled by other parameters, notaxis.Final Answer:
It decides whether to sort rows or columns in an array. -> Option AQuick Check:
axiscontrols direction = A [OK]
- Confusing axis with sorting order
- Thinking axis changes data type
- Assuming axis sets sorting algorithm
arr along rows?Solution
Step 1: Identify axis for sorting rows
In a 2D array, axis=1 means sorting each row individually.Step 2: Check other options for validity
Axis=0 sorts columns, axis=None flattens the array into 1D before sorting, axis=2 is invalid for 2D arrays.Final Answer:
numpy.sort(arr, axis=1) -> Option CQuick Check:
axis=1 sorts rows = B [OK]
- Using axis=0 to sort rows
- Using invalid axis like 2 for 2D arrays
- Using axis=None which flattens the array
import numpy as np arr = np.array([[3, 1, 2], [6, 4, 5]]) sorted_arr = np.sort(arr, axis=0) print(sorted_arr)
Solution
Step 1: Understand sorting along axis=0
Sorting with axis=0 sorts each column independently in ascending order.Step 2: Sort each column of the array
Columns: [3,6] -> [3,6], [1,4] -> [1,4], [2,5] -> [2,5]. Since columns are already sorted, array remains the same.Final Answer:
[[3 1 2] [6 4 5]] sorted by columns -> Option BQuick Check:
axis=0 sorts columns = A [OK]
- Assuming sorting rearranges rows
- Confusing axis=0 with axis=1
- Expecting full array sort instead of column-wise
import numpy as np arr = np.array([[1, 3], [2, 4]]) sorted_arr = np.sort(arr, axis=2) print(sorted_arr)
Solution
Step 1: Check array dimensions
The array is 2D with shape (2,2), so valid axes are 0 and 1 only.Step 2: Validate axis parameter
Using axis=2 is invalid and causes an IndexError because axis 2 does not exist.Final Answer:
Axis 2 does not exist for a 2D array. -> Option AQuick Check:
Axis must be within array dimensions = C [OK]
- Using axis out of range
- Assuming sort only works on 1D arrays
- Confusing axis with array shape
arr with shape (2, 2, 3), how would you sort the array along the last axis for each 2D slice?Solution
Step 1: Identify the last axis in a 3D array
For shape (2, 2, 3), axes are 0, 1, 2. The last axis is 2.Step 2: Use axis=2 to sort along the last axis
Sorting with axis=2 sorts each 2D slice along the last dimension (length 3).Final Answer:
np.sort(arr, axis=2) -> Option DQuick Check:
Last axis is 2, so axis=2 sorts last dimension [OK]
- Using axis=0 or 1 instead of last axis
- Confusing negative axis indexing
- Not matching axis to array shape
