Numpy use an array to slice the indices of another array

Question

Context: I have a 2D array A that I would like to modify at specific indices, given by the arrays a1 and a2. I could use a for loop but I want to optimize this problem.

Problem: The way I modify my array needs to use a2 as a slice: A[a1, a2:] = 1000. But I can't manage to get past TypeError: only integer scalar arrays can be converted to a scalar index. How could I do that value replacement of A faster than with loops ?

Example:

import numpy as np

# Initialize array
A = np.zeros((10,10),int)

    
# Create two arrays of indices
a1 = np.array([1,5,6], dtype = int)
a2 = np.array([4,6,2], dtype = int)

# As a for loop
for i in range(a1.shape[0]):
    A[a1[i], a2[i]:] = 10

# What I tried (doesn't work)
A[a1, a2:]

A
Out[452]: 
array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0, 10, 10, 10, 10, 10, 10],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0, 10, 10, 10, 10],
       [ 0,  0, 10, 10, 10, 10, 10, 10, 10, 10],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0]])

Since the rows and columns are being indexed over different arrays, I don't think you can avoid looping: stackoverflow.com/a/60279765/7776212 — diviquery
– diviquery, Commented Jun 26, 2023 at 23:05

hpaulj · Accepted Answer · 2023-06-26 23:50:27Z

Your example, tweaked for easier display:

In [450]: A = np.zeros((10,10),int)
     ...: # Create two arrays of indices
     ...: a1 = np.array([1,5,6], dtype = int)
     ...: a2 = np.array([4,6,2], dtype = int)
     ...: 

In [451]: for i in range(a1.shape[0]):
     ...:     A[a1[i], a2[i]:] = 10
     ...: 

In [452]: A
Out[452]: 
array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0, 10, 10, 10, 10, 10, 10],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0, 10, 10, 10, 10],
       [ 0,  0, 10, 10, 10, 10, 10, 10, 10, 10],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0]])

Ranges can be converted to advanced indexing arrays. In this case:

In [453]: idx2=np.arange(10)>=a2[:,None]
In [454]: idx2
Out[454]: 
array([[False, False, False, False,  True,  True,  True,  True,  True,
         True],
       [False, False, False, False, False, False,  True,  True,  True,
         True],
       [False, False,  True,  True,  True,  True,  True,  True,  True,
         True]])

And we can find the '10s' in A with:

In [455]: A[a1][idx2]
Out[455]: 
array([10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10])

This approach is inspired by various posts about padding lists of various lengths.

But we can't use that to set values in A. I'll have to do some more experimenting to get that to work.

edit

Define another index that covers all rows:

In [459]: a3 = np.zeros(10,int)+10; a3[a1]=a2    
In [460]: a3
Out[460]: array([10,  4, 10, 10, 10,  6,  2, 10, 10, 10])

Then make the mask:

In [461]: idx2=np.arange(10)>=a3[:,None]
In [462]: idx2
Out[462]: 
array([[False, False, False, False, False, False, False, False, False,
        False],
       [False, False, False, False,  True,  True,  True,  True,  True,
         True],
       [False, False, False, False, False, False, False, False, False,
        False],
       [False, False, False, False, False, False, False, False, False,
        False],
       [False, False, False, False, False, False, False, False, False,
        False],
       [False, False, False, False, False, False,  True,  True,  True,
         True],
       [False, False,  True,  True,  True,  True,  True,  True,  True,
         True],
       [False, False, False, False, False, False, False, False, False,
        False],
       [False, False, False, False, False, False, False, False, False,
        False],
       [False, False, False, False, False, False, False, False, False,
        False]])

Test the fetch:

In [463]: A[idx2]
Out[463]: 
array([10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10])

and the set:

In [464]: A[idx2]=20

In [465]: A
Out[465]: 
array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0, 20, 20, 20, 20, 20, 20],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0, 20, 20, 20, 20],
       [ 0,  0, 20, 20, 20, 20, 20, 20, 20, 20],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0]])

Thanks a lot for your answer ! Should I wait for you to modify your code or should I accept it as an answer and close the topic ? I wasn't expecting my question to actually touch something non-trivial !

Collectives™ on Stack Overflow

Numpy use an array to slice the indices of another array

1 Answer 1

edit

1 Comment

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

edit

1 Comment

Linked

Related