Pythonic approach to pull and group list items based on a list of indices with tuples and non-tuples
Question:
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.
Answers:
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))
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]
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']
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.
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))
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]
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']