如何用NumPy/SciPy高效修改二维数组的2×2块？

核心思路 #

原代码用Python循环逐个处理2×2块，每次都要做tobytes()、字典查找和赋值，Python循环的开销在块数量大时会非常高。NumPy的优化方向是批量提取所有块、批量映射转换、批量替换回原网格，把循环放到C层执行，大幅提升效率。

具体实现步骤 #

1. 提取所有2×2块 #

先把网格中需要处理的区域（从yshift、xshift开始，步长2的块区域）提取出来，用view_as_blocks（来自scikit-image，或手动reshape）把区域拆分成(n_blocks_y, n_blocks_x, 2, 2)的块结构，每个元素对应一个2×2块。

2. 转换块为可快速映射的键 #

原字典用bytes作为键，效率较低。我们可以把每个2×2块的4个值（假设是bool类型）转成uint8，再通过view直接转换成一个uint32整数，用整数作为字典键，查找速度更快。

3. 批量映射替换 #

把所有块的整数键批量映射到对应的新值，再把映射后的结果重塑回块结构，替换回原网格的对应区域。

完整代码示例 #

import numpy as np
from skimage.util import view_as_blocks

# 示例转换规则（和原代码的transitions对应）
transitions = {
    b'\x00\x00\x00\x00': np.array([True, True, True, True]),
    b'\x01\x01\x01\x01': np.array([False, False, False, False]),
    b'\x01\x00\x01\x00': np.array([True, False, True, False]),
}

# 把原bytes键的字典转为uint32整数键的字典，提升查找效率
transitions_int = {
    np.frombuffer(key, dtype=np.uint32)[0]: val.astype(bool)
    for key, val in transitions.items()
}

# 生成测试网格（替换成你的实际grid）
ny, nx = 10, 10
yshift, xshift = 0, 0
grid = np.random.randint(0, 2, size=(ny, nx), dtype=bool)

# 计算可处理的块数量
n_blocks_y = (ny - yshift) // 2
n_blocks_x = (nx - xshift) // 2

# 提取需要处理的块区域（确保区域大小是2的整数倍）
block_region = grid[yshift:yshift + n_blocks_y*2, xshift:xshift + n_blocks_x*2]

# 拆分成2×2的块结构，形状为(n_blocks_y, n_blocks_x, 2, 2)
blocks = view_as_blocks(block_region, block_shape=(2, 2))
# 把每个块展平为4个元素的数组，形状变为(n_blocks_y, n_blocks_x, 4)
blocks_flat = blocks.reshape(n_blocks_y, n_blocks_x, 4)

# 将块转换为uint8类型，再通过view转成uint32整数（每个块对应一个整数键）
blocks_uint8 = blocks_flat.astype(np.uint8)
block_int_keys = blocks_uint8.view(np.uint32).reshape(n_blocks_y, n_blocks_x)

# 准备映射用的键数组和值数组
keys_arr = np.array(list(transitions_int.keys()), dtype=np.uint32)
vals_arr = np.array(list(transitions_int.values()), dtype=bool)

# 批量找到每个块键对应的索引（确保所有键都在字典中）
indices = np.searchsorted(keys_arr, block_int_keys)
assert np.all(keys_arr[indices] == block_int_keys), "存在未定义的块模式"

# 获取所有块的新值
new_blocks_flat = vals_arr[indices]

# 重塑回块结构并替换原网格的对应区域
new_blocks = new_blocks_flat.reshape(n_blocks_y, n_blocks_x, 2, 2)
grid[yshift:yshift + n_blocks_y*2, xshift:xshift + n_blocks_x*2] = new_blocks.reshape(block_region.shape)

关键优化点 #

• 替换Python循环为批量操作：所有块的提取、转换、映射都在NumPy的C层执行，避免了Python循环的逐次开销，块数量越多，效率提升越明显。
• 用整数键替代bytes键：字典查找整数的速度远快于bytes，且通过view转换几乎没有额外开销。
• 利用view_as_blocks简化块处理：无需手动调整数组形状，代码更简洁易读。

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