3242. 设计相邻元素求和服务

最新推荐文章于 2025-12-20 09:26:15 发布

原创最新推荐文章于 2025-12-20 09:26:15 发布 · 288 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#python #开发语言 #leetcode #数据结构 #算法

python 专栏收录该内容

1691 篇文章

订阅专栏

利用 Python 实现二维数组中元素的邻接与对角和查询 —— NeighborSum 类详解

在实际开发中，我们经常需要在二维数组中进行位置相关的计算，比如图像处理、游戏地图、路径规划、矩阵搜索等任务中，“相邻元素”的概念极其重要。

今天我们将深入介绍一个经典的问题：给定一个二维数组，如何快速查询某个元素的相邻（上下左右）与对角（左上、右上、左下、右下）元素之和？
我们将基于这个需求，实现一个简洁高效的类 NeighborSum，并探索优化过程。

📌 问题描述

我们有一个 n x n 的二维数组 grid，里面包含 [0, n² - 1] 范围内的不重复整数。我们希望实现以下功能：

初始化时记录 grid 结构。
adjacentSum(value)：查询该值在矩阵中上下左右四个相邻元素的和。
diagonalSum(value)：查询该值在矩阵中四个对角方向的元素之和。

举个例子，如果 grid 是：

[
 [0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]
]

那么 adjacentSum(4) 应该返回 1 + 3 + 5 + 7 = 16，diagonalSum(4) 应该返回 0 + 2 + 6 + 8 = 16。

🔨 基础实现

我们先来看一份直接明了的实现，它使用了传统的方式来写逻辑判断：

from typing import List

class NeighborSum:
    def __init__(self, grid: List[List[int]]):
        self.value_map = {}
        self.grid = grid
        self.row = len(grid)
        self.col = len(grid[0])
        for i in range(self.row):
            for j in range(self.col):
                self.value_map[grid[i][j]] = (i, j)

    def adjacentSum(self, value: int) -> int:
        row, col = self.value_map[value]
        top = self.grid[row - 1][col] if row > 0 else 0
        bottom = self.grid[row + 1][col] if row < self.row - 1 else 0
        left = self.grid[row][col - 1] if col > 0 else 0
        right = self.grid[row][col + 1] if col < self.col - 1 else 0
        return top + bottom + left + right

    def diagonalSum(self, value: int) -> int:
        row, col = self.value_map[value]
        top_left = self.grid[row - 1][col - 1] if row > 0 and col > 0 else 0
        top_right = self.grid[row - 1][col + 1] if row > 0 and col < self.col - 1 else 0
        bottom_left = self.grid[row + 1][col - 1] if row < self.row - 1 and col > 0 else 0
        bottom_right = self.grid[row + 1][col + 1] if row < self.row - 1 and col < self.col - 1 else 0
        return top_left + top_right + bottom_left + bottom_right

该实现具有：

✅ 清晰的逻辑
✅ 高性能的查找（使用字典存储值到坐标的映射）
❌ 代码略显重复，可维护性不高

🚀 优化版本：使用方向数组

为了让代码更整洁，我们可以使用方向数组和 Python 的海象运算符（:=，Python 3.8+）来进一步精简逻辑：

from typing import List

class NeighborSum:
    def __init__(self, grid: List[List[int]]):
        self.grid = grid
        self.rows = len(grid)
        self.cols = len(grid[0])
        self.value_map = {
            grid[i][j]: (i, j)
            for i in range(self.rows)
            for j in range(self.cols)
        }

    def adjacentSum(self, value: int) -> int:
        row, col = self.value_map[value]
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # 上下左右
        return sum(
            self.grid[r][c]
            for dr, dc in directions
            if 0 <= (r := row + dr) < self.rows and 0 <= (c := col + dc) < self.cols
        )

    def diagonalSum(self, value: int) -> int:
        row, col = self.value_map[value]
        directions = [(-1, -1), (-1, 1), (1, -1), (1, 1)]  # 四个对角方向
        return sum(
            self.grid[r][c]
            for dr, dc in directions
            if 0 <= (r := row + dr) < self.rows and 0 <= (c := col + dc) < self.cols
        )

优化点：

✅ 使用方向数组减少冗余代码
✅ 使用字典推导式一行构造坐标映射
✅ 更具函数式编程风格，结构更优雅

🧪 测试样例

你可以用以下代码进行测试：

grid = [
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8]
]

ns = NeighborSum(grid)
print(ns.adjacentSum(4))  # 输出 16
print(ns.diagonalSum(4))  # 输出 16
print(ns.adjacentSum(0))  # 输出 4（右+下）
print(ns.diagonalSum(0))  # 输出 4（右下）