CODEWITHSUNDEEP
pythonnumpy
orange
orange

Reputation: 8068

2d index to select elements from 1d array

I'm trying to use a 2d boolean array (ix) to pick elements from a 1d array (c) to create a 2d array (r). The resulting 2d array is also a boolean array. Each column stands for the unique value in c.

Example:

>>> ix
array([[ True,  True, False, False, False,  False, False],
       [False, False,  True, False, False, False,  True],
       [False, False, False,  True, False, False, False]])
>>> c
array([1, 2, 3, 4, 8, 2, 4])

Expected result

          1,     2,     3,     4,     8
r = [
     [ True,  True, False, False, False], # c[ix[0][0]] == 1 and c[ix[0][1]] == 2; it doesn't matter that ix[0][5] (pointing to `2` in `c`) is False as ix[0][1] was already True which is sufficient.
     [False, False,  True,  True, False], # [3]
     [False, False, False,  True, False]  # [4] as ix[2][3] is True
    ]

Can this be done in a vectorised way?

Upvotes: 2

Views: 49

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150735

Let us try:

# unique values
uniques = np.unique(c)

# boolean index into each row
vals = np.tile(c,3)[ix.ravel()]

# search within the unique values
idx = np.searchsorted(uniques, vals)

# pre-populate output
out = np.full((len(ix), len(uniques)), False)

# index into the output:
out[np.repeat(np.arange(len(ix)), ix.sum(1)), idx ] = True

Output:

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

Upvotes: 2

Related Questions