python中的数组操作

Question

我们假设我有这个简单的数组:

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

如果我们将此列表视为一个表,其中列由逗号和由元组分隔的行分隔,我想创建一个只检索我想要的列的函数.例如,这个函数看起来像这样:

get_columns(array, tuple_columns_selector))

例如,我希望只收集第一列和第三列,在这种情况下,它会返回另一个带有新值的数组:

如果我做:

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

它将返回如下内容:

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

等等.你能帮我创建这个get_columns函数吗？这是我尝试过的代码:

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

使用operator.itemgetter和的魔力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)

返回：

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

cols什么时候(0, 2)。

它返回：

[1,2,3]

cols什么时候(1,)。

所以你的 get_columns 函数可以是

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