How to use np.where multiple times without iteration?

Question

Note that this question is not about multiple conditions within a single np.where(), see this thread for that.

I have a numpy array arr1 with some numbers (without a particular structure):

arr0 = \
np.array([[0,3,0],
          [1,3,2],
          [1,2,0]])

and a list of all the entries in this array:

entries = [0,1,2,3]

I also have another array, arr1:

arr1 = \
np.array([[4,5,6],
          [6,2,4],
          [3,7,9]])

I would like to perform some function on multiple subsets of elements of arr1. A subset consts of numbers which are at the same position as arr0 entries with a cetrain value. Let this function be finding the max value. Performing the function on each subset via a list comprehension:

res = [np.where(arr0==index,arr1,0).max() for index in entries]

res is [9, 6, 7, 5]

As expected: 0 in arr0 is on the top left, top right, bottom right corner, and the biggest number from the top left, top right, bottom right entries of arr1 (ie 4, 6, 9) is 9. Rest follow with a similar logic.

How can I achieve this without iteration?

My actual arrays are much bigger than these examples.

´arr0´ and ´arr1´ is 2000 by 2000, while entries have about ´100´ numbers in it. — zabop, Nov 18 '21 at 22:59

Reti43 · Accepted Answer · 2021-11-19T13:05:50.447

With broadcasting

res = np.where(arr0[...,None] == entries, arr1[...,None], 0).max(axis=(0, 1))

The result of np.where(...) is a (3, 3, 4) array, where slicing [...,0] would give you the same 3x3 array you get by manually doing the np.where with just entries[0], etc. Then taking the max of each 3x3 subarray leaves you with the desired result.

Timings

Apparently this method doesn't scale well for bigger arrays. The other answer using np.unique is more efficient because it reduces the maximum operation down to a few unique value regardless of how big the original arrays are.

import timeit

import matplotlib.pyplot as plt
import numpy as np

def loops():
    return [np.where(arr0==index,arr1,0).max() for index in entries]
def broadcast():
    return np.where(arr0[...,None] == entries, arr1[...,None], 0).max(axis=(0, 1))
def numpy_1d():
    arr0_1D = arr0.ravel()
    arr1_1D = arr1.ravel()
    arg_idx = np.argsort(arr0_1D)
    u, idx = np.unique(arr0_1D[arg_idx], return_index=True)
    return np.maximum.reduceat(arr1_1D[arg_idx], idx)

sizes =  (3, 10, 25, 50, 100, 250, 500, 1000)
lengths = (4, 10, 25, 50, 100)
methods = (loops, broadcast, numpy_1d)

fig, ax = plt.subplots(len(lengths), sharex=True)
for i, M in enumerate(lengths):
    entries  = np.arange(M)
    times = [[] for _ in range(len(methods))]
    for N in sizes:
        arr0 = np.random.randint(1000, size=(N, N))
        arr1 = np.random.randint(1000, size=(N, N))
        for j, method in enumerate(methods):
            times[j].append(np.mean(timeit.repeat(method, number=1, repeat=10)))
    for t in times:
        ax[i].plot(sizes, t)
    ax[i].legend(['loops', 'broadcasting', 'numpy_1d'])
    ax[i].set_title(f'Entries size {M}')

plt.xticks(sizes)
fig.text(0.5, 0.04, 'Array size (NxN)', ha='center')
fig.text(0.04, 0.5, 'Time (s)', va='center', rotation='vertical')
plt.show()

Great numpy usage! This however, in my actual case takes longer time than the iteration, so I'll leave this question open for a bit if others add a faster alternative. Otherwise, I'll accept this answer. — zabop, Nov 18 '21 at 16:30
@zabop For the current example broadcasting seems to be twice as fast on my machine (even though they are both blazing fast). Is it possible that your actual arrays are much bigger and broadcasting to a bigger dimension make this memory inefficient? It would be useful if you could edit your question to add that information in. — Reti43, Nov 18 '21 at 16:40

score 2 · Answer 2 · answered Nov 18 '21 at 18:35

It's more convenient to work in 1D case. You need to sort your arr0 then find starting indices for every group and use np.maximum.reduceat.

arr0_1D = np.array([[0,3,0],[1,3,2],[1,2,0]]).ravel()
arr1_1D = np.array([[4,5,6],[6,2,4],[3,7,9]]).ravel()
arg_idx = np.argsort(arr0_1D)
>>> arr0_1D[arg_idx]
array([0, 0, 0, 1, 1, 2, 2, 3, 3])
u, idx = np.unique(arr0_1D[arg_idx], return_index=True)
>>> idx
array([0, 3, 5, 7], dtype=int64)
>>> np.maximum.reduceat(arr1_1D[arg_idx], idx)
array([9, 6, 7, 5], dtype=int32)

How to use np.where multiple times without iteration?

2 Answers2

Timings