Numpy: Insert value into different column for each row

Question

I'm trying to insert into a numpy matrix given a mask that defines a single cell per row. Effectively, it's inserting a value into each row but with a different column. I've tried to use

np.insert()

without success:

>>> x

array([[False, False,  True, False, False],

       [False, False,  True, False, False],

       [False, False,  True, False, False],

       [False, False,  True, False, False],

       [False, False,  True, False, False]], dtype=bool)

>>> y = np.arange(25).reshape(5,5)

>>> y

array([[ 0,  1,  2,  3,  4],

       [ 5,  6,  7,  8,  9],

       [10, 11, 12, 13, 14],

       [15, 16, 17, 18, 19],

       [20, 21, 22, 23, 24]])

>>> np.insert(y, np.where(x)[1], 99, axis=1)

array([[ 0,  1, 99, 99, 99, 99, 99,  2,  3,  4],

       [ 5,  6, 99, 99, 99, 99, 99,  7,  8,  9],

       [10, 11, 99, 99, 99, 99, 99, 12, 13, 14],

       [15, 16, 99, 99, 99, 99, 99, 17, 18, 19],

       [20, 21, 99, 99, 99, 99, 99, 22, 23, 24]])

Anytime I try and insert based on the

x

mask, it ends up duplicating values.

Also as noted, the mask may potentially be setup in a way that it's not a simple column. For example:

>>> x = np.zeros((5, 5), dtype=bool)

>>> x[1:, 2] = True

>>> x[0, 1] = True

>>> x

array([[False,  True, False, False, False],

       [False, False,  True, False, False],

       [False, False,  True, False, False],

       [False, False,  True, False, False],

       [False, False,  True, False, False]], dtype=bool)

Which then means I can't simply specify a particular column as the index to insert at:

>>> np.insert(y, 2, [99, 99, 99, 99, 99], axis=1)

array([[ 0,  1, 99,  2,  3,  4],

       [ 5,  6, 99,  7,  8,  9],

       [10, 11, 99, 12, 13, 14],

       [15, 16, 99, 17, 18, 19],

       [20, 21, 99, 22, 23, 24]])

The desired output would be:

array([[ 0,  99, 1,  2,  3,  4],

       [ 5,  6, 99,  7,  8,  9],

       [10, 11, 99, 12, 13, 14],

       [15, 16, 99, 17, 18, 19],

       [20, 21, 99, 22, 23, 24]])

Any help would be greatly appreciated!

Recommended for you: Get network issues from WhatsUp Gold. Not end users.

Answer 1

Approach #1 : Here's one way with boolean-indexing -

def insert_one_per_row(arr, mask, putval):
    mask_ext = np.column_stack((mask, np.zeros((len(mask),1),dtype=bool)))
    out = np.empty(mask_ext.shape, dtype=arr.dtype)
    out[~mask_ext] = arr.ravel()
    out[mask_ext] = putval
    return out

Sample run -

In [88]: y
Out[88]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

In [89]: x
Out[89]: 
array([[False,  True, False, False, False],
       [False, False,  True, False, False],
       [False, False, False, False,  True],
       [ True, False, False, False, False],
       [False, False,  True, False, False]], dtype=bool)

In [90]: insert_one_per_row(y, x, putval=99)
Out[90]: 
array([[ 0, 99,  1,  2,  3,  4],
       [ 5,  6, 99,  7,  8,  9],
       [10, 11, 12, 13, 99, 14],
       [99, 15, 16, 17, 18, 19],
       [20, 21, 99, 22, 23, 24]])

We can also assign different values per row -

In [91]: insert_one_per_row(y, x, putval=[-1,-2,-3,-4,-5])
Out[91]: 
array([[ 0, -1,  1,  2,  3,  4],
       [ 5,  6, -2,  7,  8,  9],
       [10, 11, 12, 13, -3, 14],
       [-4, 15, 16, 17, 18, 19],
       [20, 21, -5, 22, 23, 24]])

Approach #2 : We will get the flattened True places on the mask and insert the new values with np.insert on a flattened version of the input array at those places, like so -

def insert_one_per_row_v2(arr, mask, putval):
    idx = np.flatnonzero(mask)
    return np.insert(arr.ravel(), idx, putval).reshape(arr.shape[0],-1)

Numpy: Insert value into different column for each row

Sign up