⚡️ Speed up function groupby_mean
by 6,392%
#46
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
📄 6,392% (63.92x) speedup for
groupby_mean
insrc/numpy_pandas/dataframe_operations.py
⏱️ Runtime :
278 milliseconds
→4.28 milliseconds
(best of338
runs)📝 Explanation and details
Here is an optimized rewrite of your program. The main bottleneck in your original code is the use of
df.iloc[i][col]
inside a Python loop, which is extremely slow (iloc
is not efficient for row-wise access in a loop, since it creates new Series each time and is pure Python). We can extract both columns as numpy arrays (fast), then use a single loop over these pre-extracted arrays, vastly reducing overhead. However, the fastest approach is to use pandas' own highly optimized groupby mechanism, which is written in C. Computing group means withgroupby().mean()
is both correct and orders of magnitude faster.I'll preserve your function signature and structure, but internally use vectorized pandas operations for speed, then convert the output to a dict as in your original output.
If it is absolutely required not to use groupby():
Here is a version that manually aggregates the data but without the per-row iloc access overhead:
Both versions will run much faster than the original code.
If maximum speed is the goal, always use the first version with
groupby().mean()
.✅ Correctness verification report:
🌀 Generated Regression Tests and Runtime
To edit these changes
git checkout codeflash/optimize-groupby_mean-mcfa14jd
and push.