NumPy grouping using itertools.groupby performance

NumPy

itertools

groupby

performance

Python

NumPy grouping using itertools.groupby performance

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Start Practicing Learn More

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:

Method	Time	Relative Speed
itertools.groupby	2.5s	1x
pandas groupby	0.15s	17x
np.bincount	0.02s	125x

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