How to Efficiently Index N-Dimensional Arrays with Lower-Dimensional Index Arrays?

Indexing an N-Dimensional Array with an (N-1)-Dimensional Array

Accessing an N-dimensional array with an (N-1)-dimensional array presents a challenge when seeking values aligned along a specific dimension. Conventional approaches using np.argmax may not be sufficient.

Advanced Indexing Approach

Elegant indexing can be achieved through advanced indexing using np.ogrid. For a 3D array a and its argmax along the first dimension, idx:

import numpy as np

a = np.random.random_sample((3, 4, 4))
idx = np.argmax(a, axis=0)

m, n = a.shape[1:]
I, J = np.ogrid[:m, :n]
a_max_values = a[idx, I, J]

This approach creates a grid that effectively expands the index array to the full dimensions of the original array.

Generalization for Arbitrary Dimensions

For a more generalized solution, the argmax_to_max() function can be defined:

def argmax_to_max(arr, argmax, axis):
    new_shape = list(arr.shape)
    del new_shape[axis]

    grid = np.ogrid[tuple(map(slice, new_shape))]
    grid.insert(axis, argmax)

    return arr[tuple(grid)]

This function takes the original array, its argmax, and the desired axis and returns the corresponding maximum values.

Alternative Approach for General Indexing

For indexing any N-dimensional array with an (N-1)-dimensional array, the all_idx() function is a more simplified solution:

def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)

Using this function, indexing into the array a with idx along axis axis can be accomplished with:

axis = 0
a_max_values = a[all_idx(idx, axis=axis)]

The above is the detailed content of How to Efficiently Index N-Dimensional Arrays with Lower-Dimensional Index Arrays?. For more information, please follow other related articles on the PHP Chinese website!