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]])

Answer 1

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]])