Pythonic approach to pull and group list items based on a list of indices with tuples and non-tuples-CodePudding

Consider a list of things. In my actual problem, the things are matplotlib artists, but for a more generalized case, lets call the list, list_of_things:

list_of_things = ["one", "two", "three", "four", "five"]

Additionally, we have a list of indices into list_of_things in which list elements may or may not be grouped into tuples. Let's call this list_of_indices:

list_of_indices = [(3, 1), (2, 0), 4]

The desired result is a new list containing items from list_of_things which preserves the order and shape of items in list_of_indices, like this:

desired_result = [("four", "two"), ("three", "one"), "five"]

One approach to solve this is through a loop, using an empty list as a collector of the results:

results = []
for item in list_of_indices:
    if isinstance(item, tuple):
        results.append(
            (list_of_things[item[0]], 
             list_of_things[item[1]])
        )
    else:
        results.append(list_of_things[item])
print(results)
>>> [('four', 'two'), ('three', 'one'), 'five']

But this feels obtuse. It seems like there should be a more pythonic, optimized way to do the same thing.

For the curious, what I am after is the ability to group matplotlib artist handles (i.e., things plotted in a pyplot.axes) together so that I can combine legend items using the matplotlib.legend_handler.HandlerTuple method. This comes in handy for cases where for example you are plotting a regression and associated confidence intervals, and only want to show one legend entry. Tupled handles are plotted together.

CodePudding user response：

operator.itemgetter would be a good choice (to fetch values by arbitrary indices):

from operator import itemgetter

res = [itemgetter(*(el if isinstance(el, tuple) else [el]))(list_of_things)
       for el in list_of_indices]

[('four', 'two'), ('three', 'one'), 'five']

CodePudding user response：

We can create map method for each of the types we find in the list of the indices. Since we have only two types, these are enough:

def map_int(lst, i):
    return lst[i]

def map_tuple(lst, tpl):
    return tuple(lst[t] for t in tpl)

Then, we should have a factory method that calls the right mapper: (Notice I used pattern matching, you can also use simple ifs)

def map_all(lst, tpl_or_int):
    match tpl_or_int:
        case int(i):
            return map_int(lst, i)
        case tuple(tpl):
            return map_tuple(lst, tpl)

And finally, we can use our factory method to map all our indices.

list_of_things = ["one", "two", "three", "four", "five"]
list_of_indices = [(3, 1), (2, 0), 4]

result = list(map(lambda i: map_all(list_of_things, i), list_of_indices))

CodePudding user response：

You can use numpy to do that:

>>> [np.array(list_of_things)[(i,)].tolist() for i in list_of_indices]

[['four', 'two'], ['three', 'one'], 'five']

Obviously if list_of_things is already a numpy array, you don't need to cast as array:

list_of_things = np.array(list_of_things)
out = [list_of_things[(i,)].tolist() for i in list_of_indices]