Practice Now
NumPy
itertools
groupby
performance
Python

NumPy grouping using itertools.groupby performance

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Data ScienceComplexity & PerformancePython

Introduction

Grouping data in NumPy arrays often leads developers to itertools.groupby, but this approach is slow because it requires sorting the array, converting to Python objects, and iterating element-by-element. NumPy-native alternatives — np.unique with return_counts/return_inverse, boolean indexing, and np.bincount — are 10-100x faster because they operate entirely in C without Python loop overhead. For DataFrame-like groupby operations, pandas is usually the best choice.

itertools.groupby with NumPy (Slow)

python
1import numpy as np
2from itertools import groupby
3
4data = np.array([3, 1, 2, 1, 3, 2, 1, 3, 2])
5values = np.array([10, 20, 30, 40, 50, 60, 70, 80, 90])
6
7# Sort by group key
8sorted_idx = np.argsort(data)
9sorted_keys = data[sorted_idx]
10sorted_vals = values[sorted_idx]
11
12# Group with itertools
13for key, group_iter in groupby(zip(sorted_keys, sorted_vals), key=lambda x: x[0]):
14    group_values = [v for _, v in group_iter]
15    print(f"Group {key}: {group_values}, sum={sum(group_values)}")
16
17# Group 1: [20, 40, 70], sum=130
18# Group 2: [30, 60, 90], sum=180
19# Group 3: [10, 50, 80], sum=140

This works but is slow: argsort is fast (O(n log n)), but groupby iterates in Python, creating Python objects for each element. For large arrays, this bottleneck dominates.

Fast Alternative: np.unique with return_inverse

python
1import numpy as np
2
3keys = np.array([3, 1, 2, 1, 3, 2, 1, 3, 2])
4values = np.array([10, 20, 30, 40, 50, 60, 70, 80, 90])
5
6# Get unique groups and indices
7unique_keys, inverse = np.unique(keys, return_inverse=True)
8
9# Sum values per group using bincount
10group_sums = np.bincount(inverse, weights=values)
11
12for key, total in zip(unique_keys, group_sums):
13    print(f"Group {key}: sum={total}")
14
15# Group 1: sum=130.0
16# Group 2: sum=180.0
17# Group 3: sum=140.0

np.unique with return_inverse maps each element to its group index. np.bincount with weights computes weighted sums per group — all in C, no Python loop.

Counting Elements per Group

python
1keys = np.array([3, 1, 2, 1, 3, 2, 1, 3, 2])
2
3# Method 1: np.unique with return_counts
4unique, counts = np.unique(keys, return_counts=True)
5print(dict(zip(unique, counts)))  # {1: 3, 2: 3, 3: 3}
6
7# Method 2: np.bincount (only for non-negative integers)
8counts = np.bincount(keys)
9print(counts)  # [0, 3, 3, 3] — index 0 has 0 occurrences

np.bincount is the fastest option for non-negative integer keys. np.unique with return_counts works with any data type.

Group Mean, Min, Max

python
1import numpy as np
2
3keys = np.array([0, 1, 0, 1, 0, 1, 0, 1])
4values = np.array([10, 20, 30, 40, 50, 60, 70, 80])
5
6unique_keys = np.unique(keys)
7
8# Group means
9group_sums = np.bincount(keys, weights=values)
10group_counts = np.bincount(keys)
11group_means = group_sums / group_counts
12print(f"Means: {dict(zip(unique_keys, group_means))}")
13# {0: 40.0, 1: 50.0}
14
15# Group max/min — use np.maximum.reduceat on sorted data
16sorted_idx = np.argsort(keys)
17sorted_vals = values[sorted_idx]
18sorted_keys = keys[sorted_idx]
19
20# Find where groups change
21split_points = np.searchsorted(sorted_keys, unique_keys)
22group_maxes = np.maximum.reduceat(sorted_vals, split_points)
23print(f"Maxes: {dict(zip(unique_keys, group_maxes))}")
24# {0: 70, 1: 80}

np.maximum.reduceat applies the maximum ufunc to segments of the sorted array. Combined with searchsorted to find group boundaries, this avoids Python loops entirely.

Boolean Indexing for Simple Groups

python
1keys = np.array(['A', 'B', 'A', 'B', 'A'])
2values = np.array([10, 20, 30, 40, 50])
3
4# Simple approach for a small number of groups
5for group in np.unique(keys):
6    mask = keys == group
7    group_vals = values[mask]
8    print(f"Group {group}: {group_vals}, mean={group_vals.mean():.1f}")
9
10# Group A: [10 30 50], mean=30.0
11# Group B: [20 40], mean=30.0

For a small number of groups (under 10-20), boolean indexing is simple and fast. Each pass creates a mask and selects matching elements — all vectorized.

pandas GroupBy (Best for Complex Operations)

python
1import pandas as pd
2import numpy as np
3
4keys = np.array([3, 1, 2, 1, 3, 2, 1, 3, 2])
5values = np.array([10, 20, 30, 40, 50, 60, 70, 80, 90])
6
7df = pd.DataFrame({'key': keys, 'value': values})
8
9# One line for any aggregation
10result = df.groupby('key')['value'].agg(['sum', 'mean', 'count', 'min', 'max'])
11print(result)
12#      sum  mean  count  min  max
13# key
14# 1    130  43.3      3   20   70
15# 2    180  60.0      3   30   90
16# 3    140  46.7      3   10   80

For complex groupby operations (multiple aggregations, string keys, custom functions), pandas is the practical choice. The overhead of creating a DataFrame is negligible compared to the flexibility gained.

Performance Comparison

python
1import timeit
2import numpy as np
3from itertools import groupby
4
5n = 1_000_000
6keys = np.random.randint(0, 100, n)
7values = np.random.rand(n)
8
9# itertools.groupby (requires sort)
10def groupby_itertools():
11    idx = np.argsort(keys)
12    sk, sv = keys[idx], values[idx]
13    return {k: sum(v for _, v in g) for k, g in groupby(zip(sk, sv), key=lambda x: x[0])}
14
15# np.bincount
16def groupby_bincount():
17    return np.bincount(keys, weights=values)
18
19# pandas
20import pandas as pd
21df = pd.DataFrame({'key': keys, 'value': values})
22def groupby_pandas():
23    return df.groupby('key')['value'].sum()
24
25t1 = timeit.timeit(groupby_itertools, number=10)
26t2 = timeit.timeit(groupby_bincount, number=10)
27t3 = timeit.timeit(groupby_pandas, number=10)
28
29print(f"itertools.groupby: {t1:.3f}s")
30print(f"np.bincount:       {t2:.3f}s")
31print(f"pandas groupby:    {t3:.3f}s")

Typical results for 1M elements, 100 groups:

MethodTimeRelative Speed
itertools.groupby2.5s1x
pandas groupby0.15s17x
np.bincount0.02s125x

Common Pitfalls

  • Forgetting to sort before itertools.groupby: itertools.groupby groups consecutive equal elements. Without sorting first, it creates multiple groups for the same key. Always argsort and reindex before using it.
  • Using itertools.groupby for large arrays: The Python-level iteration overhead makes it 50-100x slower than NumPy-native approaches. Use np.bincount, np.unique, or pandas for arrays over a few thousand elements.
  • np.bincount only works with non-negative integers: If your keys are floats, strings, or negative numbers, np.bincount does not work. Use np.unique(return_inverse=True) to map arbitrary keys to integer indices first.
  • Memory usage with np.unique on large arrays: np.unique sorts a copy of the array internally, doubling memory usage. For very large arrays, consider processing in chunks or using pandas which handles memory more efficiently.
  • Ignoring pandas for complex grouping: Pure NumPy groupby is fast for simple sum/mean/count but becomes complex for multi-column grouping, custom aggregations, or string keys. pandas handles these cases with simpler code and acceptable performance.

Summary

  • Avoid itertools.groupby for NumPy arrays — it requires sorting and Python-level iteration
  • Use np.bincount(keys, weights=values) for the fastest group sums with integer keys
  • Use np.unique(return_inverse=True) to map arbitrary keys to integer indices
  • Use np.maximum.reduceat on sorted data for group min/max operations
  • Use pandas groupby() for complex aggregations, string keys, or multi-column grouping
  • np.bincount is 50-125x faster than itertools.groupby for large arrays
Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.