如何用Numpy/Itertools计算N×N0-1矩阵中连续1的周边元素和？

问题背景 #

给定一个N×N的随机0-1矩阵，需要计算其中连续1区块（连续≥2个1）的周边元素之和。例如某矩阵第5行的连续1「111」，其周边元素和为1。目前已通过itertools.groupby实现连续1的检测及索引获取，希望用Numpy、Itertools或其他工具更高效地获取周边元素数组或直接计算其和。

已有的检测连续1及索引的代码：

from itertools import groupby

def groups(l): 
    return [sum(g) for i, g in groupby(l) if i == 1] 

con += list(filter(lambda x: x > 1, groups(matrix[4])))

idx += [idx for idx, i in enumerate(matrix[4]) if i == 1]

解决方案 #

方法一：Numpy高效处理（推荐大矩阵） #

Numpy的切片操作能大幅提升矩阵处理效率，结合itertools.groupby定位连续1区块，再通过边界判断提取周边元素求和：

import numpy as np
from itertools import groupby

# 生成N×N随机0-1矩阵（示例N=5）
N = 5
matrix = np.random.randint(0, 2, size=(N, N))
print("原矩阵：")
print(matrix)

# 定位所有连续≥2个1的区块，返回(row_idx, start_col, end_col)
def get_continuous_1_blocks(arr):
    blocks = []
    for row_idx, row in enumerate(arr):
        for val, group in groupby(enumerate(row), key=lambda x: x[1]):
            if val == 1:
                indices = [idx for idx, _ in group]
                if len(indices) >= 2:
                    blocks.append((row_idx, indices[0], indices[-1]))
    return blocks

# 计算所有连续1区块的周边元素和
def calculate_surround_sum(matrix, blocks):
    total_sum = 0
    rows, cols = matrix.shape
    for row, start_col, end_col in blocks:
        surround_elements = []
        # 上一行对应列
        if row > 0:
            surround_elements.extend(matrix[row-1, start_col:end_col+1].tolist())
        # 下一行对应列
        if row < rows - 1:
            surround_elements.extend(matrix[row+1, start_col:end_col+1].tolist())
        # 当前行左邻列
        if start_col > 0:
            surround_elements.append(matrix[row, start_col-1])
        # 当前行右邻列
        if end_col < cols - 1:
            surround_elements.append(matrix[row, end_col+1])
        block_sum = sum(surround_elements)
        total_sum += block_sum
        print(f"第{row+1}行连续1区块[{start_col}-{end_col}]的周边和：{block_sum}")
    return total_sum

# 执行计算
blocks = get_continuous_1_blocks(matrix)
total_surround_sum = calculate_surround_sum(matrix, blocks)
print(f"所有连续1区块的周边元素总和：{total_surround_sum}")

方法二：纯Itertools+列表推导（适合小矩阵） #

如果不需要Numpy依赖，可直接优化原代码，通过区间定位计算周边和：

from itertools import groupby
import numpy as np

# 生成N×N随机0-1矩阵（示例N=5）
N = 5
matrix = [[np.random.randint(0,2) for _ in range(N)] for _ in range(N)]
print("原矩阵：")
for row in matrix:
    print(row)

# 获取某行中连续≥2个1的区间(start, end)
def get_continuous_intervals(row):
    intervals = []
    current_start = None
    for idx, val in enumerate(row):
        if val == 1:
            if current_start is None:
                current_start = idx
        else:
            if current_start is not None:
                if idx - current_start >= 2:
                    intervals.append((current_start, idx-1))
                current_start = None
    # 处理行尾的连续1
    if current_start is not None and len(row) - current_start >= 2:
        intervals.append((current_start, len(row)-1))
    return intervals

# 计算指定行的连续1区块周边和
row_idx = 4  # 示例第5行（索引4）
row = matrix[row_idx]
intervals = get_continuous_intervals(row)
rows_num = len(matrix)
cols_num = len(row)

for start, end in intervals:
    surround_sum = 0
    # 上一行对应列
    if row_idx > 0:
        surround_sum += sum(matrix[row_idx-1][start:end+1])
    # 下一行对应列
    if row_idx < rows_num - 1:
        surround_sum += sum(matrix[row_idx+1][start:end+1])
    # 当前行左邻列
    if start > 0:
        surround_sum += matrix[row_idx][start-1]
    # 当前行右邻列
    if end < cols_num - 1:
        surround_sum += matrix[row_idx][end+1]
    print(f"第{row_idx+1}行连续1区块[{start}-{end}]的周边和：{surround_sum}")

关键说明 #

• 边界处理：所有操作都加入了索引越界判断，避免访问矩阵外的元素
• 效率差异：Numpy方法在大矩阵场景下比纯Python循环快数倍，小矩阵两者差异不大
• 区块定义：仅处理连续≥2个1的区块，符合原代码中filter(lambda x: x>1)的逻辑

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