从列表中提取二维数组的 Pythonic 方法

Posted 2023-02-23

【中文标题】从列表中提取二维数组的 Pythonic 方法

【英文标题】：Pythonic way to extract 2D array from list

【发布时间】：2018-01-13 07:08:11

【问题描述】：

假设我有一个包含 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' 

我想从这个一维列表中提取一个子矩阵作为另一个总是从第一个元素开始的一维列表。

例如从 lst 中提取一个 2 x 2 矩阵：

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

或者从 lst 中提取一个 3 x 3 矩阵：

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

这可行，但我想知道是否有更 Pythonic 的方式来做到这一点，因为我认为这种方法不能很好地适应矩阵大小。

【参考方案1】：

调用flatten() 然后map() 的效率低于：

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

这避免了一些副本和不必要的函数调用。

有关reshape() 与flatten() 的更多信息，请参阅：What is the difference between flatten and ravel functions in numpy?

您也可以完全不使用 NumPy。在某种程度上，这更简单。您需要使用特定的输入数据进行测试，看看哪个更快。

[lst[i*init_mat + j] for i in range(mat_size) for j in range(mat_size)]

【参考方案2】：

一种基于数组的方法是 -

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() 附加到输出的附加步骤。

---

如果你想避免将整个列表转换为数组，可能是因为元素数量太大而要提取的窗口相对较小，我们可以在@的帮助下生成块的有效索引987654321@。然后，将其索引到输入列表中，以作为列表的最终输出。因此，我们最终会得到这样的结果 -

idx = (np.arange(size)[:,None]*4 + np.arange(size)).ravel()
out = [lst[i] for i in idx]

【讨论】：

我接受了这个答案，因为它是第一个，而且在初步测试中似乎也是最快的。谢谢，也感谢其他用户的帮助。

【参考方案3】：

在没有 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)))

或者您可以使用compress从每 4 项中取出前 3 项

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