Python中使用np.fromiter生成含指定第四列的多维数组问题

我在Jupyter Notebook中尝试生成一个多维数组，期望得到指定的输出结果。我尝试了多段代码，但始终无法生成数值范围为30-35的第四列，目前最接近的代码如下：

import numpy as np
from itertools import chain

def fun(i):
    return tuple(4*i + j for j in range(4))

a = np.fromiter(chain.from_iterable(fun(i) for i in range(6)), 'i', 6 * 4)
a.shape = 6, 4

print(repr(a))

我期望的结果是：

array([[ 1,  2,  3, 30],
       [ 4,  5,  6, 31],
       [ 7,  8,  9, 32],
       [10, 11, 12, 33],
       [13, 14, 15, 34],
       [20, 21, 22, 35]])

解决方案 #

原有代码的生成逻辑和目标数组规律不匹配，以下两种修改方式可以得到期望结果：

方法一：分步构建数组 #

先拆分生成前三个列的内容，再拼接第四列，同时处理第六行的特殊值：

import numpy as np

# 生成前5行的前3列：1-15按3个一组拆分
first_part = np.arange(1, 16).reshape(5, 3)
# 添加第六行的前3列
full_first_three = np.vstack([first_part, [20, 21, 22]])
# 生成第四列并拼接成最终数组
target_array = np.hstack([full_first_three, np.arange(30, 36).reshape(-1, 1)])

print(repr(target_array))

方法二：修改生成函数 #

调整fun函数逻辑，分别处理前五行和第六行的前三个元素，直接生成第四列的目标值：

import numpy as np
from itertools import chain

def fun(i):
    # 前五行的前三个元素按连续3数递增
    if i < 5:
        first_three = tuple(3*i + 1 + j for j in range(3))
    # 第六行使用特殊值
    else:
        first_three = (20, 21, 22)
    # 第四列从30开始依次递增
    fourth = 30 + i
    return first_three + (fourth,)

a = np.fromiter(chain.from_iterable(fun(i) for i in range(6)), 'i', 6 * 4)
a.shape = 6, 4

print(repr(a))

两种方法运行后都能输出你期望的数组。

内容的提问来源于stack exchange，提问作者DanielS