CODEWITHSUNDEEP
pythonarraysnumpy
user17242583
user17242583

Reputation:

How to convert values to their index

I have a numpy array containing 1's and 0's:

a = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
              [0, 0, 0, 0, 0, 1, 0, 1, 0, 0],
              [0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
              [0, 1, 1, 0, 0, 0, 1, 0, 0, 0],
              [0, 0, 1, 0, 0, 1, 0, 0, 0, 1],
              [0, 0, 0, 0, 0, 1, 0, 0, 0, 1],
              [0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
              [0, 0, 0, 1, 0, 0, 1, 0, 1, 0],
              [1, 0, 0, 0, 0, 0, 0, 0, 0, 1],
              [0, 1, 1, 0, 0, 1, 1, 0, 0, 0]])

I'd like to convert each 1 to the index in the subarray that it's occuring at, to get this:

e = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
              [0, 0, 0, 0, 0, 5, 0, 7, 0, 0],
              [0, 0, 2, 0, 0, 0, 0, 0, 0, 0],
              [0, 1, 2, 0, 0, 0, 6, 0, 0, 0],
              [0, 0, 2, 0, 0, 5, 0, 0, 0, 9],
              [0, 0, 0, 0, 0, 5, 0, 0, 0, 9],
              [0, 0, 2, 0, 0, 0, 6, 0, 0, 0],
              [0, 0, 0, 3, 0, 0, 6, 0, 8, 0],
              [0, 0, 0, 0, 0, 0, 0, 0, 0, 9],
              [0, 1, 2, 0, 0, 5, 6, 0, 0, 0]])

So far what I've done is multiply the array by a range:

a * np.arange(a.shape[0])

which is good, but I'm wondering if there's a better, simpler way to do it, like a single function call?

Upvotes: 1

Views: 234

Answers (3)

Michael Szczesny
Michael Szczesny

Reputation: 5036

I benchmarked the obligatory np.einsum approach, which was ~1.29x slower for larger arrays (100_000, 1000) than the corrected original solution. The inplace solution was ~8x slower than np.einsum.

np.einsum('ij,j->ij', a, np.arange(a.shape[1]))

Upvotes: 1

Alain T.
Alain T.

Reputation: 42139

Your approach is a fast as it gets but it uses the wrong dimension for the multiplication (it would fait if the matrix wasn't square).

Multiply the matrix by a range of column indexes:

import numpy as np
a = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
              [0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0],
              [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
              [0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0],
              [0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0],
              [0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0],
              [0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
              [0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0],
              [1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
              [0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0]])

e = a * np.arange(a.shape[1])

print(e)
[[ 0  0  0  0  0  0  0  0  0  0  0]
 [ 0  0  0  0  0  5  0  7  0  0  0]
 [ 0  0  2  0  0  0  0  0  0  0  0]
 [ 0  1  2  0  0  0  6  0  0  0  0]
 [ 0  0  2  0  0  5  0  0  0  9  0]
 [ 0  0  0  0  0  5  0  0  0  9  0]
 [ 0  0  2  0  0  0  6  0  0  0 10]
 [ 0  0  0  3  0  0  6  0  8  0  0]
 [ 0  0  0  0  0  0  0  0  0  9  0]
 [ 0  1  2  0  0  5  6  0  0  0  0]]

Upvotes: 2

Warren Weckesser
Warren Weckesser

Reputation: 114956

This modifies a in place:

In [4]: i, j = np.nonzero(a)

In [5]: a[i, j] = j

In [6]: a
Out[6]: 
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 5, 0, 7, 0, 0],
       [0, 0, 2, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 2, 0, 0, 0, 6, 0, 0, 0],
       [0, 0, 2, 0, 0, 5, 0, 0, 0, 9],
       [0, 0, 0, 0, 0, 5, 0, 0, 0, 9],
       [0, 0, 2, 0, 0, 0, 6, 0, 0, 0],
       [0, 0, 0, 3, 0, 0, 6, 0, 8, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 9],
       [0, 1, 2, 0, 0, 5, 6, 0, 0, 0]])

Make a copy if you don't want modify a in place.

Or, this creates a new array (in one line):

In [8]: np.arange(a.shape[1])[a]
Out[8]: 
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 5, 0, 7, 0, 0],
       [0, 0, 2, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 2, 0, 0, 0, 6, 0, 0, 0],
       [0, 0, 2, 0, 0, 5, 0, 0, 0, 9],
       [0, 0, 0, 0, 0, 5, 0, 0, 0, 9],
       [0, 0, 2, 0, 0, 0, 6, 0, 0, 0],
       [0, 0, 0, 3, 0, 0, 6, 0, 8, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 9],
       [0, 1, 2, 0, 0, 5, 6, 0, 0, 0]])

Upvotes: 2

Related Questions