numpy
vectorization
arrays
Python
performance optimization

Vectorization of a function dependent on 2 arrays in numpy

Master System Design with Codemia

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

Introduction

When a NumPy function depends on two arrays, the fastest solution is usually to express the relationship in array operations instead of iterating in Python. In practice, that means learning how broadcasting, element-wise operations, and reduction functions let NumPy apply the same computation over whole grids of values at C speed.

Start by Identifying the Shape Relationship

The first question is whether the function combines the arrays element by element or produces every pairwise combination between them.

Element-wise example:

python
1import numpy as np
2
3a = np.array([1.0, 2.0, 3.0])
4b = np.array([10.0, 20.0, 30.0])
5
6result = a * b + 1
7print(result)

Here a[i] only interacts with b[i].

Pairwise example:

python
1import numpy as np
2
3x = np.array([1.0, 2.0, 3.0])
4y = np.array([10.0, 20.0])
5
6result = x[:, None] + y[None, :]
7print(result)

Here every value of x interacts with every value of y. That shape distinction determines how you vectorize the function.

Element-Wise Vectorization

If your original loop looks like "for each index, compute something from a[i] and b[i]," then NumPy element-wise expressions are usually enough.

Suppose the scalar function is:

f(a, b) = a^2 + 3b

Vectorized version:

python
1import numpy as np
2
3a = np.array([1.0, 2.0, 3.0])
4b = np.array([4.0, 5.0, 6.0])
5
6result = a ** 2 + 3 * b
7print(result)

This is real vectorization because NumPy performs the arithmetic on the full arrays without a Python loop.

Pairwise Vectorization With Broadcasting

A more interesting case is when the function depends on all pairs from two arrays. Broadcasting is the core NumPy tool for this.

Suppose the scalar function is:

f(x, y) = (x - y)^2

The vectorized pairwise form is:

python
1import numpy as np
2
3x = np.array([1.0, 2.0, 3.0])
4y = np.array([10.0, 20.0])
5
6result = (x[:, None] - y[None, :]) ** 2
7print(result)

x[:, None] turns x into a column vector, and y[None, :] turns y into a row vector. NumPy then broadcasts them into a 2D grid automatically.

This is the pattern to remember for many "two arrays, all combinations" problems.

Do Not Confuse np.vectorize With True Speedups

np.vectorize can make code look cleaner, but it usually does not provide the performance benefit people expect. It still calls the Python function repeatedly under the hood.

python
1import numpy as np
2
3def f(a, b):
4    return a ** 2 + 3 * b
5
6vf = np.vectorize(f)
7
8a = np.array([1.0, 2.0, 3.0])
9b = np.array([4.0, 5.0, 6.0])
10
11print(vf(a, b))

This can be useful for convenience, but it is not the same as rewriting the function with native NumPy operations. When performance is the goal, prefer broadcasting and built-in array math.

Reduce After Vectorizing

Many real tasks compute a pairwise matrix and then reduce it, for example taking a minimum or a sum.

python
1import numpy as np
2
3x = np.array([1.0, 2.0, 3.0])
4y = np.array([10.0, 20.0])
5
6distances = np.abs(x[:, None] - y[None, :])
7closest_per_x = distances.min(axis=1)
8
9print(closest_per_x)

This pattern is common in nearest-neighbor style calculations, similarity scoring, and cost-matrix construction.

Common Pitfalls

  • Not deciding whether the problem is element-wise or pairwise before writing code.
  • Using np.vectorize and expecting major performance gains.
  • Forgetting to reshape with None or np.newaxis when broadcasting is required.
  • Creating a huge pairwise matrix when the array sizes make that memory usage impractical.
  • Writing a Python loop around NumPy operations and losing most of the performance benefit.

Summary

  • The right vectorization strategy depends on whether the two arrays interact element by element or in all pairwise combinations.
  • Element-wise problems usually become direct NumPy arithmetic expressions.
  • Pairwise problems are often solved with broadcasting using shapes such as x[:, None] and y[None, :].
  • 'np.vectorize is mostly a convenience wrapper, not a true performance optimization.'
  • Think about both computation cost and memory cost before building large pairwise arrays.

Course illustration
Course illustration

All Rights Reserved.