Извлечь переменное количество столбцов для матрицы

Question

Я хочу извлечь некоторые указанные столбцы из матрицы.Моя матрица

matrix=[[1,2,3,4,5],
        [6,7,8,9,10],
        [6,4,3,1,2],
        [2,3,4,5,6]]
expected result: [[3,4],[8,9],[3,1],[4,5]] for 2 columns
expected result in case of 3 column:[[1,2,3],[6,7,8],[6,4,3],[2,3,4]]

Я пытаюсь использовать метод, указанный ниже:

def func(matrix, p):
     return np.vstack([a[i] for a in matrix])

Вышеуказанный метод просто возвращается в один столбец, однако я хочу написать метод, которыйпринимает в качестве входных данных несколько столбцов (например, ожидаемые результаты для 2 и 3 столбцов), также мое входное число столбцов меняется каждый разПожалуйста, предложите подходящий метод для извлечения этих столбцов в Python.

Answer 1

Используя чистый Python, вы можете извлечь столбцы 1, 5 и 7, используя следующее понимание вложенного списка:

[[a[i] for i in (1, 5, 7)] for a in matrix]

Пример использования для ввода:

>>> [[a[i] for i in (0, 1, 2)] for a in matrix]
[[1, 2, 3], [6, 7, 8], [6, 4, 3], [2, 3, 4]]

В подобных случаях, когда вы хотите извлечь реальный фрагмент столбцов, вы также можете использовать нарезку списка:

>>> [a[0:3] for a in matrix]
[[1, 2, 3], [6, 7, 8], [6, 4, 3], [2, 3, 4]]

Answer 2

Если вы предпочитаете работать в Numpy, вы можете использовать нарезку:

def col_slice(arr, col_start=None, total_cols=None):
    """ take a slice of columns from col_start for total_cols """
    return arr[:,col_start:col_start + total_cols]

Примеры:

In [22]: col_slice(your_matrix, col_start=2, total_cols=2)
Out[22]:
array([[3, 4],
       [8, 9],
       [3, 1],
       [4, 5]])

In [23]: col_slice(your_matrix, col_start=0, total_cols=3)
Out[23]:
array([[1, 2, 3],
       [6, 7, 8],
       [6, 4, 3],
       [2, 3, 4]])