Practice Now
Numpy
Python
Array Processing
Function Mapping
Data Manipulation

Most efficient way to map function over numpy array

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 ScienceData StructuresPython

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
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.