Using list of lists of indices to slice columns and obtain the row-wise vector length

Question

I have an NxM array, as well as an arbitrary list of sets of column indices I'd like to use to slice the array. For example, the 3x3 array

my_arr = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])

and index sets

my_idxs = [[0, 1], [2]]

I would like to use the pairs of indices to select the corresponding columns from the array and obtain the length of the (row-wise) vectors using np.linalg.norm(). I would like to do this for all index pairs. Given the aforementioned array and list of index sets, this should give:

[[2.23606797749979, 3],
 [2.23606797749979, 3],
 [2.23606797749979, 3]]

When all sets have the same number of indices (for example, using my_idxs = [[0, 1], [1, 2]] I can simply use np.linalg.norm(my_arr[:, my_idxs], axis=1):

[[2.23606797749979, 3.605551275463989],
 [2.23606797749979, 3.605551275463989],
 [2.23606797749979, 3.605551275463989]]

However, when they are not (as is the case with my_idxs = [[0, 1], [2]], the varying index list lengths yield an error when slicing as the array of index sets would be irregular in shape. Is there any way to implement the single-line option, without resorting to looping over the list of index sets and handling each of them separately?

Answer 1

You can try:

my_arr = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
my_idxs = [[0, 1], [2]]

out = np.c_[*[np.linalg.norm(my_arr[:, i], axis=1) for i in my_idxs]]
print(out)

Prints:

[[2.23606798 3.        ]
 [2.23606798 3.        ]
 [2.23606798 3.        ]]