Pythonic способ извлечь 2D-массив из списка

Скажем, у меня есть список, содержащий 16 элементов:

lst=['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P']

Этот список представляет собой массив 4 x 4, в котором все элементы помещены в одномерный список. В виде массива он имеет следующий вид:

'A', 'B', 'C', 'D'
'E', 'F', 'G', 'H'
'I', 'J', 'K', 'L'
'M', 'N', 'O', 'P' 

Я хочу извлечь подматрицу из этого списка 1D как еще один список 1D, который всегда начинается с первого элемента.

например Извлечение матрицы 2 x 2 из lst:

'A', 'B', 'E', 'F'

Или извлечение матрицы 3 x 3 из lst:

'A', 'B', 'C', 'E', 'F', 'G', 'I', 'J', 'K'

Для этого я использую numpy, чтобы изменить размер списка в массив, извлечь подматрицу, а затем снова сгладить:

import numpy as np

# The size of the matrix represented by lst
init_mat = 4
# Desired matrix size to extract
mat_size = 2
A = np.resize(lst,(init_mat,init_mat))
B = A[0:mat_size, 0:mat_size].flatten()
C = map(str,B)

Это работает, но мне было интересно, есть ли более питонический способ сделать это, так как я не думаю, что этот метод будет хорошо масштабироваться с размером матрицы.


person Yeti    schedule 05.08.2017    source источник


Ответы (3)


Один из подходов на основе массива:

size = 2 # or 3 or any number <= 4
np.asarray(lst).reshape(4,4)[:size,:size].ravel()

Пробный запуск -

In [55]: lst=['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P']

In [56]: size=2

In [57]: np.asarray(lst).reshape(4,4)[:size,:size].ravel()
Out[57]: 
array(['A', 'B', 'E', 'F'],
      dtype='|S1')

In [58]: size=3

In [59]: np.asarray(lst).reshape(4,4)[:size,:size].ravel()
Out[59]: 
array(['A', 'B', 'C', 'E', 'F', 'G', 'I', 'J', 'K'],
      dtype='|S1')

Если вам нужен массив 2D, пропустите часть ravel().

Если вы хотите получить список в качестве вывода, нам нужен дополнительный шаг .tolist(), добавляемый к выводу.


Если вы хотите избежать преобразования всего списка в массив, возможно, из-за того, что количество элементов слишком велико, а извлекаемое окно относительно меньше, мы можем просто сгенерировать допустимые индексы для блока с помощью NumPy broadcasting. Затем проиндексируйте его во входном списке для окончательного вывода в виде списка. Таким образом, мы получим что-то вроде этого -

idx = (np.arange(size)[:,None]*4 + np.arange(size)).ravel()
out = [lst[i] for i in idx]
person Divakar    schedule 05.08.2017
comment
Я принял этот ответ, так как он был первым, а также оказался самым быстрым при предварительном тестировании. Спасибо, и спасибо другим пользователям за их помощь. - person Yeti; 06.08.2017

Вызов flatten(), затем map() менее эффективен, чем:

B = A[:mat_size, :mat_size].reshape(-1)
C = B.tolist()

Это позволяет избежать некоторых копий и ненужных вызовов функций.

Подробнее о reshape() и flatten() см.: Что разница между функциями flatten и ravel в numpy?

Вы также можете сделать это вообще без NumPy. В каком-то смысле это проще. Вам нужно будет протестировать ваши конкретные входные данные, чтобы увидеть, что быстрее.

[lst[i*init_mat + j] for i in range(mat_size) for j in range(mat_size)]
person John Zwinck    schedule 05.08.2017

Делая это без numpy и учитывая, что матрица становится большой, я бы использовал итератор для обхода списка, чтобы во время извлечения не создавались дополнительные списки. Используя islice для извлечения необходимых элементов, он вырезал нужные элементы при каждой операции нарезки. В случае извлечения матрицы 3x3 первый срез начнется с индекса 0 и остановится перед индексом 3, таким образом, первые три элемента будут удалены из итератора. Следующие срезы будут начинаться с индекса 1, потому что 4 - 3 = 1, и останавливаться перед 4.

from itertools import chain, islice, repeat

lst=['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P']

width = 4
extract = [3, 3]

slice_starts = chain([0], repeat(width - extract[0]))
slice_stops = chain([extract[0]], repeat(width))
rows = map(islice, repeat(iter(lst), extract[1]), slice_starts, slice_stops)
print(list(chain.from_iterable(rows)))

Или вы можете взять первые три элемента из каждых 4 предметов, используя compress

from itertools import chain, compress, repeat

lst=['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P']

width = 4
extract = [3, 3]
selectors = repeat([i < extract[0] for i in range(width)], extract[1])
print(list(compress(lst, chain.from_iterable(selectors))))
person frogcoder    schedule 05.08.2017