Best/Fast way to remove a column from a Python 2D list

Question

I have a list of list of lists (all of lists have same size) in python like this:

A = [[1,2,3,4],['a','b','c','d'] , [12,13,14,15]]

I want to remove some columns (i-th elements of all lists).

Is there any way that does this without for statements?

Answer 1

As mentioned, you can't do this without loop. However, using built-in functions here's a functional approach that doesn't explicitly use any loop:

In [24]: from operator import itemgetter

In [25]: def remove_col(arr, ith):
    ...:     itg = itemgetter(*filter((ith).__ne__, range(len(arr[0]))))
    ...:     return list(map(list, map(itg, arr)))
    ...: 

Demo:

In [26]: remove_col(A, 1)
Out[26]: [[1, 3, 4], ['a', 'c', 'd'], [12, 14, 15]]

In [27]: remove_col(A, 3)
Out[27]: [[1, 2, 3], ['a', 'b', 'c'], [12, 13, 14]]

Note that instead of list(map(list, map(itg, arr))) if you only return map(itg, arr) it will give you the expected result but as an iterator of iterators instead of list of lists. This will be a more optimized approach in terms of both memory and run-time in this case.

Also, using loops here's the way I'd do this:

In [31]: def remove_col(arr, ith):
    ...:     return [[j for i,j in enumerate(sub) if i != ith] for sub in arr]

Surprisingly (not if you believe in the power of C :)) the functional approach is even faster for large arrays.

In [41]: arr = A * 10000

In [42]: %timeit remove_col_functional(arr, 2)
8.42 ms ± 37.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [43]: %timeit remove_col_list_com(arr, 2)
23.7 ms ± 165 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# And if in functional approach you just return map(itg, arr)
In [47]: %timeit remove_col_functional_iterator(arr, 2)
1.48 µs ± 4.71 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)