Array manipulation in python

Question

Let's suppose I have this simple array:

simple_list = [ 
 ('1', 'a', 'aa'), 
 ('2', 'b', 'bb'), 
 ('3', 'c', 'cc')
]

If we consider this list as a table, where columns are separated by comas and lines separated by tuples, I want to create a function that retrieves only the columns I want. for example, this function would look like something like this:

get_columns(array, tuple_columns_selector))

I want, for example, to collect only the first and third column out of it, in this case, it would return me another array with the new values:

if I do:

get_columns(simple_list, (0,2)) 
get_columns(simple_list, (0,))

it will return something like:

[('1', 'aa'), ('2', 'bb'), ('1', 'cc')] 
[1, 2, 3]

And so on. Could you help me creating this get_columns function, please? Here's the code I've tried:

def get_columns(arr, columns): 
 result_list = [] 
 for ii in arr: 
 for i in columns: 
 result_list.append(ii[i]) 
 return result_list 
to_do_list = [
 ('Wake Up', True), 
 ('Brush Teeh', True), 
 ('Go to work', True), 
 ('Take a shower', True), 
 ('Go to bed', False) 
] 
print(get_columns(to_do_list, (0,)))

• Chris: I show you my code if you show me yours :-) Even if its incorrect, just show what you've tried. We can help you fix it.
• Antimony: Welcome to StackOverflow. Please read How to Ask, and include details of what you have tried, specifically, show us some code that you may have tried to write.
• juanpa.arrivillaga: These are not arrays. But actually, this sounds like a great use-case for actual arrays, specifically, structured numpy.arrays
• francisco sollima: Shouldn't get_columns(simple_list, (0,)) return [('1'), ('2'), ('3')]? Since it should return a list of tuples...
• RagingRoosevelt: OP, it might be worth using a pandas dataframe. This sort of operation has built-in support if you're using data frames. It's quite a rabbit hole but totally worth it, IMO.

Answer 1

Use the magic of operator.itemgetter and map:

from operator import itemgetter
simple_list = [
 ('1', 'a', 'aa'),
 ('2', 'b', 'bb'),
 ('3', 'c', 'cc')
]
cols = (1,) # can be (0, 2)
fn = itemgetter(*cols)
print map(fn, simple_list)

Returns:

[('1', 'aa'), ('2', 'bb'), ('3', 'cc')]

when cols is (0, 2).

And it returns:

[1,2,3]

when cols is (1,).

So your get_columns function can be

def get_columns(data, cols):
 return map(itemgetter(*cols), data)

• Chris: Ah, very nice solution, much better than what I was originally thinking of. This handles the OP's logic of getting single elements or a tuple of elements very elegantly.
• Felipe Marcelus Mantovani: Cheers, I'm new to programming, I am not very familiarized to modules but that's a very interesting way of solving it. Thanks once again.

Answer 2

The answer of @kopos looks fine, I just wanted to share one without additional libraries.

simple_list = [
 ('1', 'a', 'aa'),
 ('2', 'b', 'bb'),
 ('3', 'c', 'cc')
]
def get_columns(array, tuple_columns_selector):
 return [tuple(elem[i] for i in tuple_columns_selector) for elem in array]
def get_columns_multiple_lines(array, tuple_columns_selector):
 # The only difference between the result of this version and the other is that this one returns a list of lists
 # while the other returns a list of tuples
 resulting_list = [] # Create the variable that will store the resulting list
 for elem in array: # Loop for each element in array
 resulting_list.append([]) # We add a new "empty row" to store all the columns needed
 for i in tuple_columns_selector: # Loop for each column needed
 resulting_list[-1].append(elem[i]) # We append the column value to the last item in resulting_list
 return resulting_list
print get_columns(simple_list, (0,2)) # outputs [('1', 'aa'), ('2', 'bb'), ('3', 'cc')]
print get_columns(simple_list, (0,)) # outputs [('1',), ('2',), ('3',)]
print get_columns_multiple_lines(simple_list, (0,2)) # outputs [['1', 'aa'], ['2', 'bb'], ['3', 'cc']]

The only difference is the return value when tuple_columns_selector is only one column. If it's an important difference, I can "correct" it, but you should think about how that value will be used and if it's convenient for it to have different possible structures.

• Felipe Marcelus Mantovani: Thanks a mil, you've done exactly what I wanted. It's a bit confusing for me this one lined for though. I've not yet got used to it but I will, for sure. Nah, no need to change the (1,) to 1. It's fine.
• francisco sollima: I added a more simple, not one-line approach.