GroupBy with custom functions in Pandas - Time & Space Complexity
We want to understand how the time needed changes when using groupby with custom functions in pandas.
How does the work grow as the data gets bigger?
Analyze the time complexity of the following code snippet.
import pandas as pd
def custom_func(group):
return group.sum() + group.mean()
df = pd.DataFrame({
'key': ['A', 'B', 'A', 'B', 'C'],
'value': [1, 2, 3, 4, 5]
})
result = df.groupby('key')['value'].apply(custom_func)
This code groups data by 'key' and applies a custom function to each group.
Identify the loops, recursion, array traversals that repeat.
- Primary operation: Applying the custom function to each group.
- How many times: Once per group, and inside the function, operations run over the group size.
As the total data size grows, the number of groups and group sizes affect the work.
| Input Size (n) | Approx. Operations |
|---|---|
| 10 | About 10 operations over groups and their sizes |
| 100 | About 100 operations, split across groups |
| 1000 | About 1000 operations, split across groups |
Pattern observation: The total work grows roughly in proportion to the total data size.
Time Complexity: O(n)
This means the time grows roughly in direct proportion to the number of rows in the data.
[X] Wrong: "Grouping always makes the operation slower by a square factor."
[OK] Correct: Grouping splits data, but the total work still adds up to about the size of the data, not its square.
Understanding how groupby with custom functions scales helps you explain data processing choices clearly and confidently.
"What if the custom function itself has a loop over the entire original data instead of just the group? How would the time complexity change?"
Practice
groupby function in pandas primarily do?Solution
Step 1: Understand the purpose of groupby
Thegroupbyfunction divides data into groups based on unique values in one or more columns.Step 2: Recognize the use of groups
After grouping, you can apply functions to each group separately for analysis or aggregation.Final Answer:
It splits data into groups based on column values for separate analysis. -> Option AQuick Check:
Group data = split by column values [OK]
- Confusing groupby with sorting
- Thinking groupby deletes duplicates
- Mixing groupby with merge operations
my_func to each group after grouping by column 'Category' in a DataFrame df?Solution
Step 1: Recall pandas groupby syntax
The correct method to group by a column isdf.groupby('ColumnName').Step 2: Apply a custom function to groups
Use.apply(function)after grouping to run the function on each group.Final Answer:
df.groupby('Category').apply(my_func) -> Option AQuick Check:
groupby + apply(custom function) = df.groupby('Category').apply(my_func) [OK]
- Using wrong method name like group_by
- Calling function directly on groupby object
- Misplacing apply before groupby
df below:import pandas as pd
data = {'Team': ['A', 'A', 'B', 'B', 'C'], 'Points': [10, 15, 10, 20, 30]}
df = pd.DataFrame(data)
# Custom function to get range of points in group
def point_range(x):
return x.max() - x.min()
result = df.groupby('Team')['Points'].apply(point_range).to_dict()What is the value of
result?Solution
Step 1: Group points by 'Team'
Teams: A has points [10, 15], B has [10, 20], C has [30].Step 2: Calculate range (max - min) for each group
For A: 15 - 10 = 5; B: 20 - 10 = 10; C: 30 - 30 = 0.Final Answer:
{'A': 5, 'B': 10, 'C': 0} -> Option CQuick Check:
Range = max - min per group = {'A': 5, 'B': 10, 'C': 0} [OK]
- Confusing sum with range
- Forgetting single value group range is zero
- Mixing keys and values in result
import pandas as pd
data = {'Category': ['X', 'X', 'Y'], 'Value': [1, 2, 3]}
df = pd.DataFrame(data)
# Custom function to sum values
def sum_values(x):
return x.sum
result = df.groupby('Category')['Value'].apply(sum_values)What is the issue with this code?
Solution
Step 1: Check the custom function definition
The function returnsx.sumwhich is a method, not called (missing parentheses).Step 2: Understand the effect on apply
Returning the method itself causes the result to be a method object, not the sum value.Final Answer:
The function sum_values returns the method, not the sum result. -> Option DQuick Check:
Missing () calls method, not result [OK]
- Forgetting parentheses on sum()
- Thinking groupby syntax is wrong here
- Believing apply can't use custom functions
df with columns 'Store', 'Product', and 'Sales'. You want to find the difference between the highest and lowest sales per store, but only for products with sales above 100. Which approach correctly uses groupby with a custom function to achieve this?Solution
Step 1: Understand filtering sales above 100
One approach filters inside theapplylambda on each store's sales series. Another filters the DataFrame before grouping. The third filters inside the groupapplyon the subgroup DataFrame.Step 2: Check if all approaches calculate max-min correctly
All three methods correctly compute the difference of max and min sales above 100 per store.Step 3: Confirm equivalence of results
Filtering before or inside groupby leads to the same final result for this case.Final Answer:
All of the above produce the same correct result. -> Option BQuick Check:
Filtering + groupby + custom function = all correct [OK]
- Filtering after groupby without apply
- Not filtering sales > 100 properly
- Assuming only one method works
Ready to prove it?
Pass the test and this pattern moves unknown
