Access non-sequential multiple entries in a list of lists

user_6607572 0
Viewed 60

I have a matrix of M vectors where each vector is of size N (NxM). I also have a Boolean vector of size L>=M, with exactly M entries = True. I want to create a list of lists and place the M vectors where the Boolean vector is True in same order as they are in the matrix, and the rest I want to be empty lists

Example: M = 3, N = 4, L = 5

mat = np.array([[1, 5, 9],
               [2, 6, 10],
               [3, 7, 11],
               [4, 8, 12]])
mask = [True, False, True, True, False]

I want to create the following:

res = [ [1, 2, 3, 4], [], [5, 6, 7, 8], [9, 10, 11, 12], []]

Accessing it can be done using:

data = [res[idx] for idx in range(len(res)) if mask(idx)]

However, creating it is a bit problematic. I tried creating a list of empty lists, but I can't access all relevant entries at once. Is there an elegant way of doing it?

indexing python-3.x numpy
3 Answers

Here is how I would do it:

mi = iter(mat.T.tolist())
[(m or []) and next(mi) for m in mask]
# [[1, 2, 3, 4], [], [5, 6, 7, 8], [9, 10, 11, 12], []]
user_72073920

As you are already using a list comprehension to get the data back from res, I would do a similar thing to create res in the first place.

mask_cs = np.cumsum(mask) - 1  # array([0, 0, 1, 2, 2]) , gives the corresponding index in mat
res = [mat[:, mask_cs[idx]].tolist() if mask[idx] else [] for idx in range(L)]

As alternativ, which accesses all columns of mat at once, on can create an intermediate array with size [N, L]

import numpy as np
res = np.zeros((N, L))  # Create result array
res[:, mask] = mat      # Copy the data at the right positions
res = res.T.tolist()    # Transform the array to a list of lists
for idx in range(L):    # Replace the columns with empty lists, if mask[idx] is False
    if not mask[idx]:
        res[idx] = []
user_75708170

We could make use of np.split for some elegance, like so -

In [162]: split_cols = np.split(mat.T,np.cumsum(mask)[:-1])

In [163]: split_cols
Out[163]: 
[array([[1, 2, 3, 4]]),
 array([], shape=(0, 4), dtype=int64),
 array([[5, 6, 7, 8]]),
 array([[ 9, 10, 11, 12]]),
 array([], shape=(0, 4), dtype=int64)]

So, that gives us a list of 2D arrays. For the desired output of list of lists, we need to map them to such -

In [164]: list(map(list,(map(np.ravel,split_cols))))
Out[164]: [[1, 2, 3, 4], [], [5, 6, 7, 8], [9, 10, 11, 12], []]

Alternatively, we can use lambda if that looks more elegant to some -

In [165]: F = lambda a: np.ravel(a).tolist()

In [166]: list(map(F,split_cols))
Out[166]: [[1, 2, 3, 4], [], [5, 6, 7, 8], [9, 10, 11, 12], []]
user_32938810
Related