Numpy - how to convert an array of vector indices to a mask?

Question

Given a np.ndarray named indices with a n rows and variable length vector in each row I want to create a boolean mask of n rows and m rows where m is a pre-known value equal to the greatest value possible in indices . Take note that the indices specified in indices refer to per-row indices, and not global matrix indices.

For example, given:

indices = np.array([
    [2, 0],
    [0],
    [4, 7, 1]
])

# Expected output
print(mask)
[[ True False  True False False False False False]
 [ True False False False False False False False]
 [False  True False False  True False False  True]]

m is known beforehand (the maximum length of each row in mask ) and doesn't need to be inferred from indices

Notice : This is different from converting an array of indices to a mask where the indices refer to the resulting matrix indices

Answer 1

Here's one way -

def mask_from_indices(indices, ncols=None):
    # Extract column indices
    col_idx = np.concatenate(indices)

    # If number of cols is not given, infer it based on max column index
    if ncols is None:
        ncols = col_idx.max()+1

    # Length of indices, to be used as no. of rows in o/p
    n = len(indices)

    # Initialize o/p array
    out = np.zeros((n,ncols), dtype=bool)

    # Lengths of each index element that represents each group of col indices
    lens = np.array(list(map(len,indices)))

    # Use np.repeat to generate all row indices
    row_idx = np.repeat(np.arange(len(lens)),lens)

    # Finally use row, col indices to set True values
    out[row_idx,col_idx] = 1
    return out    

Sample run -

In [89]: mask_from_indices(indices)
Out[89]: 
array([[ True, False,  True, False, False, False, False, False],
       [ True, False, False, False, False, False, False, False],
       [False,  True, False, False,  True, False, False,  True]])

Answer 2

Here is a variant:

def create_mask(indices, m):
    mask = np.zeros((len(indices), m), dtype=bool)
    for i, idx in enumerate(indices):
        mask[i, idx] = True
    return mask

Usage:

>>> create_mask(indices, 8)
array([[ True, False,  True, False, False, False, False, False],
       [ True, False, False, False, False, False, False, False],
       [False,  True, False, False,  True, False, False,  True]])

Answer 3

While there is no direct way of doing this in a fully vectorized way, for larger inputs, a single application of mask[full_row_indices, full_col_indices] with the pre-computed full list of indices is faster than multiple applications of mask[partial_row_indices, partial_col_indices] . Memory-wise, the multiple applications are also less demanding because no intermediate full_row_indices / full_col_indices need to be built. Of course this would generally depend on the length of indices .

Just to get some feeling on how much faster the different possible solutions could, the following functions have been tested:

import numpy as np
import random


def gen_mask_direct(col_indices, cols=None):
    if cols is None:
        cols = np.max(np.concatenate(col_indices)) + 1
    rows = len(col_indices)
    mask = np.zeros((rows, cols), dtype=bool)
    for row_index, col_index in enumerate(col_indices):
        mask[row_index, col_index] = True
    return mask 


def gen_mask_loops(col_indices, cols=None):
    rows = len(col_indices)
    row_indices = tuple(i for i, j in enumerate(col_indices) for _ in j)
    col_indices = tuple(sum(col_indices, ()))
    if cols is None:
        cols = np.max(col_indices) + 1
    mask = np.zeros((rows, cols), dtype=bool)
    mask[row_indices, col_indices] = True
    return mask


def gen_mask_np_repeat(col_indices, cols=None):
    rows = len(col_indices)
    lengths = list(map(len, col_indices))
    row_indices = np.repeat(np.arange(rows), lengths)
    col_indices = np.concatenate(col_indices)
    if cols is None:
        cols = np.max(col_indices) + 1
    mask = np.zeros((rows, cols), dtype=bool)
    mask[row_indices, col_indices] = True
    return mask


def gen_mask_np_concatenate(col_indices, cols=None):
    rows = len(col_indices)
    row_indices = tuple(np.full(len(col_index), i) for i, col_index in enumerate(col_indices))
    row_indices = np.concatenate(row_indices)
    col_indices = np.concatenate(col_indices)
    if cols is None:
        cols = np.max(col_indices) + 1
    mask = np.zeros((rows, cols), dtype=bool)
    mask[row_indices, col_indices] = True
    return mask

gen_mask_direct() is basically @Derlin answer and implements multiple applications of mask[partial_row_indices, partial_col_indices] . All the others implement a single application of mask[full_row_indices, full_col_indices] with different ways of preparing the full_row_indices and the full_col_indices :

• gen_mask_loops() uses direct looping
• gen_mask_np_repeat() uses np.repeat() (and it is substantially the same as @Divakar answer
• gen_mask_np_concatenate() uses a combination of np.full() and np.concatenate()

---

A quick sanity check indicates that all these are equivalent:

funcs = gen_mask_direct, gen_mask_loops, gen_mask_np_repeat, gen_mask_np_concatenate

random.seed(0)
test_inputs = [
    (tuple(
        tuple(sorted(set([random.randint(0, n - 1) for _ in range(random.randint(1, n - 1))])))
                for _ in range(random.randint(1, n - 1))))
    for n in range(5, 6)
    ]
print(test_inputs)
# [((0, 2, 3, 4), (2, 3, 4), (1, 4), (0, 1, 4))]

for func in funcs:
    print('Func:', func.__name__)
    for test_input in test_inputs:    
        print(func(test_input).astype(int))

Func: gen_mask_direct
[[1 0 1 1 1]
 [0 0 1 1 1]
 [0 1 0 0 1]
 [1 1 0 0 1]]
Func: gen_mask_loops
[[1 0 1 1 1]
 [0 0 1 1 1]
 [0 1 0 0 1]
 [1 1 0 0 1]]
Func: gen_mask_np_repeat
[[1 0 1 1 1]
 [0 0 1 1 1]
 [0 1 0 0 1]
 [1 1 0 0 1]]
Func: gen_mask_np_concatenate
[[1 0 1 1 1]
 [0 0 1 1 1]
 [0 1 0 0 1]
 [1 1 0 0 1]]

---

Here are some benchmarks (using the code from here ):

and zooming to the fastest:

supporting the overall statement that, typically, a single application of mask[...] for full indices is faster multiple applications of mask[...] for partial indices.

---

For completeness, the following code was used to generate the inputs, compare the outputs, run the benchmarks and prepare the plots:

def gen_input(n):
    random.seed(0)
    return tuple(
        tuple(sorted(set([random.randint(0, n - 1) for _ in range(random.randint(n // 2, n - 1))])))
        for _ in range(random.randint(n // 2, n - 1)))


def equal_output(a, b):
    return np.all(a == b)


input_sizes = tuple(int(2 ** (2 + (3 * i) / 4)) for i in range(13))
print('Input Sizes:\n', input_sizes, '\n')


runtimes, input_sizes, labels, results = benchmark(
    funcs, gen_input=gen_input, equal_output=equal_output,
    input_sizes=input_sizes)


plot_benchmarks(runtimes, input_sizes, labels, units='ms')
plot_benchmarks(runtimes, input_sizes, labels, units='ms', zoom_fastest=2)