np.choose() for conditional selection in NumPy - Time & Space Complexity
We want to understand how the time taken by np.choose() changes as the input size grows.
Specifically, how does the work increase when we select elements conditionally from multiple arrays?
Analyze the time complexity of the following code snippet.
import numpy as np
choices = np.array([10, 20, 30])
index = np.array([0, 2, 1, 0, 2])
result = np.choose(index, [choices, choices*2, choices*3])
This code selects elements from three arrays based on the index array, creating a new array with chosen values.
Identify the loops, recursion, array traversals that repeat.
- Primary operation: Selecting elements from multiple arrays using the
indexarray. - How many times: Once for each element in the
indexarray.
As the size of the index array grows, the number of selections grows at the same rate.
| Input Size (n) | Approx. Operations |
|---|---|
| 10 | About 10 element selections |
| 100 | About 100 element selections |
| 1000 | About 1000 element selections |
Pattern observation: The work grows linearly with the number of elements to select.
Time Complexity: O(n)
This means the time taken grows directly in proportion to the number of elements you select.
[X] Wrong: "The time grows with the number of choice arrays, so more arrays means much slower."
[OK] Correct: The main cost depends on how many elements you select, not how many arrays you have. Adding more arrays adds a small fixed cost per element, but the growth is still mainly linear with the number of elements.
Understanding how array operations scale helps you write efficient data code and explain your choices clearly in interviews.
What if we replaced np.choose() with a loop that selects elements one by one? How would the time complexity change?
Practice
np.choose() function do in NumPy?Solution
Step 1: Understand the purpose of np.choose()
np.choose()uses an index array to pick elements from multiple arrays at corresponding positions.Step 2: Compare with other options
Sorting, mean calculation, and random number generation are unrelated tonp.choose().Final Answer:
Selects elements from multiple arrays based on an index array -> Option BQuick Check:
np.choose() = Select elements [OK]
- Confusing np.choose() with sorting functions
- Thinking np.choose() calculates statistics
- Assuming np.choose() generates random values
np.choose() with an index array idx and option arrays a, b, c?Solution
Step 1: Recall np.choose() parameter order
The first argument is the index array, the second is a list of arrays to choose from.Step 2: Match syntax with options
Onlynp.choose(idx, [a, b, c])matches the correct order and structure.Final Answer:
np.choose(idx, [a, b, c]) -> Option AQuick Check:
Index first, then list of arrays [OK]
- Swapping index array and options
- Passing options as separate arguments instead of a list
- Using incorrect argument order
import numpy as np idx = np.array([0, 2, 1, 0]) a = np.array([10, 20, 30, 40]) b = np.array([100, 200, 300, 400]) c = np.array([1000, 2000, 3000, 4000]) result = np.choose(idx, [a, b, c]) print(result)
Solution
Step 1: Understand index array and options
Index array is [0, 2, 1, 0]. Arrays are a=[10,20,30,40], b=[100,200,300,400], c=[1000,2000,3000,4000].Step 2: Select elements using np.choose()
For each position, pick from arrays based on index:
pos0: idx=0 -> a[0]=10
pos1: idx=2 -> c[1]=2000
pos2: idx=1 -> b[2]=300
pos3: idx=0 -> a[3]=40
So result = [a[0], c[1], b[2], a[3]] = [10, 2000, 300, 40]Step 3: Check options
[10 300 200 40] incorrectly uses 300 and 200 instead of 2000 and 300. [10 2000 300 40] matches our result.Final Answer:
[10 2000 300 40] -> Option AQuick Check:
np.choose picks arrays by index, then elements by position [OK]
- Mixing up index meaning (array vs element)
- Swapping elements from wrong arrays
- Confusing np.choose with fancy indexing
import numpy as np idx = np.array([0, 1, 2]) a = np.array([1, 2]) b = np.array([3, 4]) c = np.array([5, 6]) result = np.choose(idx, [a, b, c]) print(result)
Solution
Step 1: Check index array and options
Index array is [0,1,2], which is valid for 3 options (a,b,c).Step 2: Check shapes of option arrays
a, b, c each have shape (2,), but index array has length 3. np.choose expects index and options to have compatible shapes.Step 3: Identify cause of error
Index array length 3 vs option arrays length 2 causes shape mismatch error.Final Answer:
Index array and option arrays have incompatible shapes causing mismatch -> Option CQuick Check:
Shape mismatch causes np.choose error [OK]
- Assuming index values cause error
- Thinking np.choose needs boolean index
- Believing np.choose limits number of options
test1 = np.array([80, 90, 70]) test2 = np.array([85, 88, 75]) test3 = np.array([78, 92, 72])And an index array indicating the best test for each student:
best = np.array([1, 2, 0])Using
np.choose(), how do you create an array of the best scores for each student?Solution
Step 1: Understand index array meaning
Index arraybest= [1, 2, 0] means:
For student 0, best test is index 1 -> test2
For student 1, best test is index 2 -> test3
For student 2, best test is index 0 -> test1Step 2: Match option arrays order to index
To pick correctly, option arrays must be in order [test1, test2, test3] so index 0 picks test1, index 1 picks test2, index 2 picks test3.Step 3: Verify np.choose usage
Usingnp.choose(best, [test1, test2, test3])picks correct scores per student.Final Answer:
np.choose(best, [test1, test2, test3]) -> Option DQuick Check:
Order options to match index meaning [OK]
- Assuming option arrays order matches test numbers
- Not matching index values to option array order
- Using wrong order causing incorrect selection
Ready to prove it?
Pass the test and this pattern moves unknown
