broadcast views irregularly numpy

Based on your example, you can simply factor the index mapping through the computation to the very end:

output2 = np.dot(unique_values * other_array1, other_array2).squeeze()[index_mapping]assert (output == output2).all()

Your expression admits two significant optimizations:

• do the indexing last
• multiply other_array1 with other_array2 first and then with unique_values

Let's apply these optimizations:

>>> output_pp = (unique_values @ (other_array1.ravel() * other_array2.ravel()))[index_mapping]# check for correctness>>> (output == output_pp).all()True# and compare it to @Yakym Pirozhenko's approach>>> from timeit import timeit>>> print("yp:", timeit("np.dot(unique_values * other_array1, other_array2).squeeze()[index_mapping]", globals=globals()))yp: 3.9105667411349714>>> print("pp:", timeit("(unique_values @ (other_array1.ravel() * other_array2.ravel()))[index_mapping]", globals=globals()))pp: 2.2684884609188884

These optimizations are easy to spot if we observe two things:

(1) if A is an mxn-matrix and b is an n-vector then

A * b == A @ diag(b)A.T * b[:, None] == diag(b) @ A.T

(2) if A is anmxn-matrix and I is a k-vector of integers from range(m) then

A[I] == onehot(I) @ A

onehot can be defined as

def onehot(I, m, dtype=int):    out = np.zeros((I.size, m), dtype=dtype)    out[np.arange(I.size), I] = 1    return out

Using these facts and abbreviating uv, im, oa1 and oa2 we can write

uv[im] * oa1 @ oa2 == onehot(im) @ uv @ diag(oa1) @ oa2

The above optimizations are now simply a matter of choosing the best order for these matrix multiplications which is

onehot(im) @ (uv @ (diag(oa1) @ oa2))

Using (1) and (2) backwards on this we obtain the optimized expression from the beginning of this post.