Python 3：获取转换为对角线列表的原矩阵索引（单词搜索场景）

我来帮你搞定对角线方向的单词查找，结合你现有的思路，我们可以从两个角度完善这个功能，先说说最直接的优化方案，再给你一个更简洁的通用实现思路。

一、基于你现有对角线思路的改进 #

你当前的做法是把对角线转成字符串搜索，核心问题是缺少字符到原矩阵坐标的映射，所以找到单词后没法反推位置。调整方法很简单：

1. 生成对角线时保存坐标信息 #

不要只存字符组成的字符串，而是把每个对角线做成包含(字符, 行号, 列号)的列表。比如针对左上到右下的对角线：

# 假设原矩阵是matrix，rows是行数，cols是列数
diagonals = []
# 按row-col的差值分组，同一差值的字符属于同一条左上到右下的对角线
for diff in range(-(cols-1), rows):
    diag = []
    for r in range(rows):
        c = r - diff
        if 0 <= c < cols:
            diag.append( (matrix[r][c], r, c) )
    diagonals.append(diag)

这样每条对角线里的元素都带着原坐标，后续找到单词时直接就能拿到位置。

2. 搜索单词并记录坐标 #

遍历每条对角线的字符序列，当找到匹配的单词时，把对应位置的坐标全部存入一个集合：

marked = set()
target_words = ["cat", "big"]

for diag in diagonals:
    # 把对角线的字符拼接成字符串用于搜索
    diag_str = ''.join([char for char, _, _ in diag])
    for word in target_words:
        word_len = len(word)
        start_idx = diag_str.find(word)
        while start_idx != -1:
            # 取出单词每个字符对应的坐标
            for i in range(word_len):
                _, r, c = diag[start_idx + i]
                marked.add( (r, c) )
            # 继续找下一个匹配
            start_idx = diag_str.find(word, start_idx + 1)

3. 生成标记后的输出 #

初始化一个全为.的矩阵，把标记集合里的坐标替换成原字符即可：

output = []
for r in range(rows):
    row = []
    for c in range(cols):
        if (r, c) in marked:
            row.append(matrix[r][c])
        else:
            row.append('.')
    output.append(' '.join(row))

print('\n'.join(output))

二、更简洁的通用方向遍历方案 #

其实不用提前生成对角线，直接在原矩阵上遍历所有可能的方向（包括对角线）会更高效，代码也更易维护。核心思路是定义方向步长，逐个检查每个起始点在对应方向上是否匹配单词：

1. 定义所有搜索方向 #

把水平、垂直、对角线的所有可能方向用步长(行增量, 列增量)表示：

directions = [
    (0, 1),   # 水平向右
    (0, -1),  # 水平向左
    (1, 0),   # 垂直向下
    (-1, 0),  # 垂直向上
    (1, 1),   # 左上→右下
    (-1, -1), # 右下→左上
    (1, -1),  # 右上→左下
    (-1, 1)   # 左下→右上
]

2. 遍历所有起始点和方向 #

对每个起始坐标，检查每个方向上是否能匹配目标单词：

marked = set()
rows = len(matrix)
cols = len(matrix[0]) if rows else 0
target_words = ["cat", "big"]

for word in target_words:
    word_len = len(word)
    if word_len == 0:
        continue
    # 遍历所有可能的起始点
    for r in range(rows):
        for c in range(cols):
            # 尝试每个方向
            for dr, dc in directions:
                # 计算单词最后一个字符的坐标，判断是否在矩阵范围内
                end_r = r + (word_len - 1) * dr
                end_c = c + (word_len - 1) * dc
                if 0 <= end_r < rows and 0 <= end_c < cols:
                    # 检查每个字符是否匹配
                    match = True
                    for i in range(word_len):
                        curr_r = r + i * dr
                        curr_c = c + i * dc
                        if matrix[curr_r][curr_c] != word[i]:
                            match = False
                            break
                    if match:
                        # 匹配成功，记录所有坐标
                        for i in range(word_len):
                            marked.add( (r + i * dr, c + i * dc) )

3. 生成最终输出 #

和前面的方法一样，用标记集合生成带.的矩阵即可。

这种方法的优势是扩展性极强，后续要加任何方向（比如反方向），只要在directions里加对应的步长就行，不需要额外修改逻辑，也省去了对角线生成的步骤。

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