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,)))

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)

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.