Most efficient way to map function over numpy array

Numpy

Python

Array Processing

Function Mapping

Data Manipulation

Most efficient way to map function over numpy array

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

The most efficient way to apply a function over a NumPy array is vectorization — using NumPy's built-in operations that execute in compiled C code. For element-wise math, use NumPy ufuncs directly (np.sqrt(arr), arr * 2). For custom functions, try to rewrite them using NumPy operations. As a last resort, use np.vectorize() or np.frompyfunc(), which are convenience wrappers around Python loops and offer no real speedup. Avoid Python for loops over NumPy arrays — they are 10-100x slower than vectorized operations.

Vectorized Operations (Fastest)

python

1import numpy as np
2
3arr = np.array([1, 4, 9, 16, 25])
4
5# Built-in ufuncs — fastest possible
6result = np.sqrt(arr)       # [1, 2, 3, 4, 5]
7result = np.sin(arr)        # Element-wise sine
8result = np.log(arr)        # Element-wise log
9result = arr ** 2           # Element-wise square
10result = arr * 2 + 1        # Element-wise arithmetic
11
12# Boolean operations
13mask = arr > 10             # [False, False, False, True, True]
14filtered = arr[mask]        # [16, 25]

NumPy ufuncs are implemented in C and operate on entire arrays without Python loop overhead. Always prefer these when possible.

np.where for Conditional Logic

python

1arr = np.array([-3, -1, 0, 2, 5])
2
3# Vectorized if-else
4result = np.where(arr > 0, arr * 2, 0)
5print(result)  # [0, 0, 0, 4, 10]
6
7# Nested conditions
8result = np.where(arr > 0, 'positive',
9         np.where(arr < 0, 'negative', 'zero'))
10print(result)  # ['negative', 'negative', 'zero', 'positive', 'positive']
11
12# Clip values to a range
13result = np.clip(arr, 0, 10)
14print(result)  # [0, 0, 0, 2, 5]

np.where is the vectorized equivalent of if-else and runs at C speed.

Rewriting Custom Functions as Vectorized

python

1# SLOW: Python function with if/else
2def categorize(x):
3    if x < 0:
4        return -1
5    elif x == 0:
6        return 0
7    else:
8        return 1
9
10# Can't directly apply to array: categorize(arr) fails
11
12# FAST: Rewrite using NumPy operations
13def categorize_vectorized(arr):
14    return np.sign(arr).astype(int)
15
16arr = np.array([-5, 0, 3, -1, 7])
17print(categorize_vectorized(arr))  # [-1, 0, 1, -1, 1]

python

1# SLOW: Custom distance function
2def custom_transform(x):
3    if x > 100:
4        return x * 0.9
5    elif x > 50:
6        return x * 0.95
7    else:
8        return x
9
10# FAST: Vectorized with np.where
11def custom_transform_fast(arr):
12    return np.where(arr > 100, arr * 0.9,
13           np.where(arr > 50, arr * 0.95, arr))
14
15arr = np.random.randint(0, 200, size=1000000)
16result = custom_transform_fast(arr)  # ~2ms vs ~800ms with loop

np.vectorize (Convenience, Not Speed)

python

1import numpy as np
2
3# np.vectorize wraps a Python function to accept arrays
4def my_func(x):
5    if x < 0:
6        return 0
7    return x ** 0.5
8
9vfunc = np.vectorize(my_func)
10arr = np.array([-4, -1, 0, 4, 9])
11result = vfunc(arr)
12print(result)  # [0, 0, 0, 2, 3]

np.vectorize does NOT make the function run faster — it is essentially a for loop with broadcasting support. It is only a convenience for making scalar functions accept arrays.

Performance Comparison

python

1import numpy as np
2import timeit
3
4arr = np.random.randn(1000000)
5
6# Method 1: Vectorized NumPy — fastest
7def vectorized():
8    return np.where(arr > 0, arr * 2, arr * -1)
9# ~3ms
10
11# Method 2: np.vectorize — slow (Python loop in disguise)
12scalar_func = lambda x: x * 2 if x > 0 else x * -1
13vfunc = np.vectorize(scalar_func)
14def with_vectorize():
15    return vfunc(arr)
16# ~400ms
17
18# Method 3: List comprehension — slowest
19def with_loop():
20    return np.array([x * 2 if x > 0 else x * -1 for x in arr])
21# ~800ms
22
23# Method 4: map() — also slow
24def with_map():
25    return np.fromiter(map(scalar_func, arr), dtype=float, count=len(arr))
26# ~500ms

Vectorized NumPy is 100-200x faster than Python loops for large arrays.

Using np.frompyfunc

python

1import numpy as np
2
3# frompyfunc is slightly faster than vectorize
4def custom(x):
5    return x ** 2 + 1
6
7ufunc = np.frompyfunc(custom, 1, 1)
8result = ufunc(np.array([1, 2, 3, 4]))
9print(result)  # [2, 5, 10, 17] — returns object array
10
11# Convert to numeric
12result = result.astype(float)

np.frompyfunc creates a ufunc that returns an object array. It is marginally faster than np.vectorize but still a Python-level loop.

Pandas apply vs NumPy

python

1import pandas as pd
2import numpy as np
3
4arr = np.random.randn(1000000)
5series = pd.Series(arr)
6
7# NumPy vectorized — fastest
8result = np.where(arr > 0, arr, 0)  # ~3ms
9
10# Pandas vectorized — fast
11result = series.where(series > 0, 0)  # ~5ms
12
13# Pandas apply — slow (Python loop)
14result = series.apply(lambda x: x if x > 0 else 0)  # ~300ms

If your data is in a pandas DataFrame, still prefer NumPy operations on the underlying .values array for best performance.

Lookup Table for Discrete Mapping

python

1# For mapping integer values to other values, use indexing
2arr = np.array([0, 1, 2, 3, 1, 0, 2])
3lookup = np.array([10, 20, 30, 40])  # 0→10, 1→20, 2→30, 3→40
4
5result = lookup[arr]
6print(result)  # [10, 20, 30, 40, 20, 10, 30]
7# This is O(1) per element — fastest possible mapping

Common Pitfalls

Using np.vectorize expecting C-speed: np.vectorize is a Python loop with a nice API. It does not compile or optimize the function. For real speedups, rewrite using NumPy operations or use Numba (@numba.vectorize).
Python for loop over large arrays: Iterating over a million-element array in Python takes hundreds of milliseconds. The same operation vectorized takes single-digit milliseconds. Always try vectorization first.
Applying pandas apply unnecessarily: df['col'].apply(lambda x: x * 2) is much slower than df['col'] * 2. Use vectorized pandas/NumPy operations before reaching for apply.
Creating intermediate arrays: np.where(arr > 0, arr, 0) creates a boolean mask array internally. For very large arrays, this doubles memory usage. Use in-place operations (np.clip(arr, 0, None, out=arr)) when memory is constrained.
Ignoring Numba for complex custom functions: When a function cannot be expressed in NumPy operations, @numba.jit compiles Python to machine code and achieves near-C performance: @numba.vectorize def f(x): return x**2 + 1.

Summary

Use NumPy built-in operations (np.sqrt, np.where, arithmetic) for maximum speed
Rewrite custom functions using NumPy vectorized operations instead of if/else
np.vectorize is for convenience, not performance — it wraps a Python loop
Avoid Python for loops and list(map()) on NumPy arrays — they are 100x slower
Use lookup table indexing (lookup[arr]) for mapping discrete integer values
For functions that cannot be vectorized, use Numba @jit for compiled performance