numpy
vectorization
arrays
Python
performance optimization

Vectorization of a function dependent on 2 arrays in numpy

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

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.

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

All Rights Reserved.