2

So the challenge is this; given an awkward array with n rows and a list of n indices (i_1 to i_n), return a list containing element i_m of row_m for all rows.

This could be done like;

import awkward
some_awkward_array = awkward.fromiter([[1,2],[3,4,5],[6]]) 
some_indices = [0,2,0]
desired_elements = [row[i] for row, i in zip(some_awkward_array, some_indices)]
assert desired_elements == [1,5,6]

But if this were a numpy array we would have access to choose and so we could do;

import numpy as np
some_numpy_array = np.array([[1,2,0],[3,4,5],[6,0,0]]) 
some_indices = [0,2,0]
desired_elements = np.choose(some_indices, some_numpy_array.T)
assert desired_elements == [1,5,6]

The second version seems to scale better, it becomes faster somewhere round 12 rows. Is there an equivalent option for an awkward array?

Edit; maybe this is an IndexedMaskedArray thing, but I can't get it to do what I want.

Clumsy cat
  • 289
  • 1
  • 12
  • 1
    I guess the first solution with the list comprehension is as good as you're going to do; the numpy version with a rectangular array is able to be faster because rectangular arrays (stored contiguously in memory) are inherently more efficient to index in two dimensions. – kaya3 Nov 07 '19 at 14:42

0 Answers0