如何将stdin读取到二维整数数组中？

Question

我想在Python中从标准输入（也就是命令行）或者文件里读取一个二维整数数组。

下面是一个不工作的代码：

from StringIO import StringIO
from array import array

# fake stdin
stdin = StringIO("""1 2
3 4
5 6""")

a = array('i')
a.fromstring(stdin.read())

运行这个代码时，我遇到了一个错误：a.fromstring(stdin.read())，错误信息是：ValueError: 字符串长度不是项大小的倍数。

Answer 1

arr = []
arr = raw_input()

arr = []
arr = raw_input().split()

如果你想根据空格来分割输入内容：

Answer 2

我之前从没用过array.array，所以我得查了查资料。

答案就在错误信息里 -

ValueError: 字符串长度不是项大小的倍数

那么，如何确定项的大小呢？这要看你初始化时用的类型。在你的例子中，你用的是i，也就是有符号整数。那么，整数有多大呢？你可以问问你的Python解释器。

>>> a.itemsize
4

上面的值给我们提供了问题的线索。你的字符串只有11个字节长，而11不是4的倍数。不过，增加字符串的长度也不会让你得到一个包含{1,2,3,4,5,6}的数组……我不太确定这样会得到什么。为什么会有这种不确定性呢？好吧，看看下面的文档说明……（现在有点晚了，所以我把重要的部分标出来了，以防你像我一样困！）

array.fromfile(f, n) 从文件对象f中读取n个项（作为机器值）并将它们添加到数组的末尾。如果可用的项少于n个，会抛出EOFError，但可用的项仍会被插入到数组中。f必须是真正的内置文件对象；其他有read()方法的东西是没用的。

array.fromstring的读取方式和array.fromfile一样。注意上面的加粗部分。“作为机器值”意味着“以二进制形式读取”。所以，要实现你想做的事情，你需要使用struct模块。看看下面的代码。

import struct
a = array.array('i')
binary_string = struct.pack('iiii', 1, 2, 3, 4)
a.fromstring(binary_string)

上面的代码片段加载了数组，里面的值是1, 2, 3, 4；正如我们预期的那样。

希望这能帮到你。

Answer 3

有几种方法可以实现这个目标。下面是一些可能的选项。

使用 数组

从列表中获取

把问题中最后一行代码替换成下面这行。

a.fromlist([int(val) for val in stdin.read().split()])

现在：

>>> a
array('i', [1, 2, 3, 4, 5, 6])

缺点：不保留二维结构（见评论）。

从生成器获取

注意：这个选项是根据eryksun的评论整合的。

更高效的方法是使用生成器，而不是列表。把问题中最后两行代码替换成：

a = array('i', (int(val) for row in stdin for val in row.split()))

这样做的结果和上面的选项一样，但避免了创建中间列表。

使用NumPy数组

如果你想保留二维结构，可以使用NumPy数组。下面是完整的例子：

from StringIO import StringIO
import numpy as np

# fake stdin
stdin = StringIO("""1 2
3 4
5 6""")

a = np.loadtxt(stdin, dtype=np.int)

现在：

>>> a
array([[1, 2],
       [3, 4],
       [5, 6]])

使用标准列表

问题中并不清楚是否可以使用Python列表。如果可以的话，一种实现目标的方法是把问题中最后两行代码替换成下面的。

a = [map(int, row.split()) for row in stdin]

运行后，我们得到了：

>>> a
[[1, 2], [3, 4], [5, 6]]