Kruskal算法实现优化求助：Borg迷宫问题Python代码超时

我正在完成Kattis平台的Borg迷宫问题，代码能正确解决问题但运行超时。已经参考他人方案优化了输入处理，让BFS找到所有外星节点后提前终止，但仍未解决超时问题。我认为排序环节影响不大，怀疑Kruskal算法部分还有优化空间，但不知从何入手，请求Python层面的优化建议。

以下是我的代码：

test_cases = int(input())
mazes = []
nodes_in_mazes = []
number_of_nodes_in_maze = []
for i in range(test_cases):
    maze = []
    rows = int(input().split()[1])
    start = ()
    aliens = []
    for r in range(rows):
        input_row = input()
        if 'S' in input_row:
            start = (r, input_row.index('S'))
        if 'A' in input_row:
            for char_index, char in enumerate(input_row):
                if char == 'A':
                    aliens.append((r, char_index))
        maze.append(input_row)
    mazes.append(maze)
    nodes_found = [start]
    for alien in aliens:
        nodes_found.append(alien)
    nodes_in_mazes.append(nodes_found)
    number_of_nodes_in_maze.append(len(nodes_found))

offsets = [[-1, 0], [0, -1], [1, 0], [0, 1]]
w_and_d_for_mazes = []


def bfs(origin, maze_nbr, alien_nodes):
    queue = [(origin[0], origin[1])]
    weight_and_distance = []
    visited_nodes = [(origin[0], origin[1])]
    in_queue = len(queue)
    weight = 1
    alien_counter = 0

    while not len(queue) == 0:
        current_square = queue.pop(0)
        for off in offsets:
            i = current_square[0] + off[0]
            j = current_square[1] + off[1]
            square = mazes[maze_nbr][i][j]
            if square != '#' and (i, j) not in visited_nodes:
                queue.append((i, j))
                visited_nodes.append((i, j))
                if (i, j) in alien_nodes:
                    weight_and_distance.append((weight, nodes_in_mazes[maze_nbr].index((origin[0], origin[1])),
                                                nodes_in_mazes[maze_nbr].index((i, j))))
                    alien_counter += 1
                    if alien_counter == number_of_nodes_in_maze[maze_nbr] - 1:
                        return weight_and_distance
        in_queue -= 1
        if in_queue == 0:
            weight += 1
            in_queue = len(queue)
    return weight_and_distance


for i in range(len(nodes_in_mazes)):
    alien_nodes = nodes_in_mazes[i].copy()
    alien_nodes.pop(0)
    maze_w_and_d = []
    for tupp in bfs(nodes_in_mazes[i][0], i, alien_nodes):
        maze_w_and_d.append(tupp)
    aliens_to_be_found = len(alien_nodes)
    while len(alien_nodes) != 0:
        for tupp in bfs(alien_nodes.pop(0), i, alien_nodes):
            maze_w_and_d.append(tupp)
    maze_w_and_d.sort()
    w_and_d_for_mazes.append(maze_w_and_d)


def find(p, i_find):
    if p[i_find] != i_find:
        p[i_find] = find(p, p[i_find])
    return p[i_find]


def union(p, ra, x, y):
    if ra[x] < ra[y]:
        p[x] = y
    elif ra[x] > ra[y]:
        p[y] = x
    else:
        p[y] = x
        ra[x] += 1


def kruskal(nodes, length):
    result = []
    ind = 0
    e = 0

    p = []
    rank = []
    minimum = 0

    for node in range(number_of_nodes_in_maze[length]):
        p.append(node)
        rank.append(0)

    while e < len(p) - 1:

        w, u, v = nodes[ind]
        ind = ind + 1
        x = find(p, u)
        y = find(p, v)

        if x != y:
            e = e + 1
            result.append([w, u, v])
            union(p, rank, x, y)
            minimum += w

    print(minimum)


for index, m in enumerate(w_and_d_for_mazes):
    kruskal(m, index)

Python层面优化建议 #

一、BFS模块核心优化（最影响性能的部分） #

1. 替换队列实现：
   目前用list做队列，pop(0)是O(n)复杂度，每次弹出都会移动整个列表元素。换成collections.deque的popleft()方法，这是O(1)操作，能大幅降低BFS的时间开销。
   示例修改：

   from collections import deque
   # BFS函数内
   queue = deque([(origin[0], origin[1])])
   # 弹出时用
   current_square = queue.popleft()
   

2. 替换visited存储结构：
   用list存储visited_nodes，(i,j) in visited_nodes是O(n)查询。换成set存储，集合的成员查询是O(1)，能减少大量重复查询的时间。
   示例修改：

   visited_nodes = set()
   visited_nodes.add((origin[0], origin[1]))
   # 判断时
   if square != '#' and (i, j) not in visited_nodes:
       visited_nodes.add((i, j))
       queue.append((i, j))
   

3. 预存节点ID映射：
   每次调用nodes_in_mazes[maze_nbr].index((x,y))是O(n)遍历查找。提前为每个迷宫建立坐标到节点ID的字典，比如pos_to_id = {pos: idx for idx, pos in enumerate(nodes_in_mazes[maze_nbr])}，之后直接通过字典取值，O(1)复杂度。
   示例修改：
   在处理每个迷宫时添加：

   pos_to_id_list = []
   # ... 处理完nodes_found后
   pos_to_id = {pos: idx for idx, pos in enumerate(nodes_found)}
   pos_to_id_list.append(pos_to_id)
   

   BFS里替换index查询：

   weight_and_distance.append((weight, pos_to_id_list[maze_nbr][origin], pos_to_id_list[maze_nbr][(i,j)]))
   

4. 简化BFS的距离计算：
   目前用in_queue计数来维护距离，逻辑复杂且容易出错。改成在队列中存储(x, y, distance)，每次扩展节点时直接传递距离+1，更直观且效率相当。
   示例修改：

   def bfs(origin, maze_nbr, alien_nodes, pos_to_id):
       queue = deque([(origin[0], origin[1], 0)])
       weight_and_distance = []
       visited_nodes = set()
       visited_nodes.add((origin[0], origin[1]))
       alien_set = set(alien_nodes)
       alien_counter = 0
       target_count = number_of_nodes_in_maze[maze_nbr] - 1
   
       while queue:
           x, y, dist = queue.popleft()
           for off in offsets:
               nx = x + off[0]
               ny = y + off[1]
               if 0 <= nx < len(mazes[maze_nbr]) and 0 <= ny < len(mazes[maze_nbr][0]):
                   square = mazes[maze_nbr][nx][ny]
                   if square != '#' and (nx, ny) not in visited_nodes:
                       visited_nodes.add((nx, ny))
                       new_dist = dist + 1
                       if (nx, ny) in alien_set:
                           weight_and_distance.append((new_dist, pos_to_id[origin], pos_to_id[(nx, ny)]))
                           alien_counter += 1
                           if alien_counter == target_count:
                               return weight_and_distance
                       queue.append((nx, ny, new_dist))
       return weight_and_distance
   

二、Kruskal算法优化 #

1. 简化初始化逻辑：
   初始化父节点列表p和秩列表rank时，不用循环append，直接用生成式：

   node_count = number_of_nodes_in_maze[length]
   p = list(range(node_count))
   rank = [0] * node_count
   

   比循环append更高效简洁。

2. 移除无用变量：
   result列表只是存储选中的边，但最终只用到minimum求和，完全可以删除result相关的所有操作，减少内存占用和不必要的赋值。

3. 迭代实现Find函数：
   递归版本的find在Python中有栈开销，换成迭代版本能避免递归深度限制且稍快：

   def find(p, i_find):
       while p[i_find] != i_find:
           p[i_find] = p[p[i_find]]  # 路径压缩
           i_find = p[i_find]
       return i_find
   

三、输入处理优化 #

多次调用input()会有IO开销，尤其是测试用例多的时候。改用一次性读取所有输入：

import sys

data = sys.stdin.read().splitlines()
ptr = 0
test_cases = int(data[ptr])
ptr +=1
mazes = []
nodes_in_mazes = []
number_of_nodes_in_maze = []
pos_to_id_list = []  # 新增存储每个迷宫的坐标到ID映射

for i in range(test_cases):
    cols, rows = map(int, data[ptr].split())
    ptr +=1
    maze = []
    start = None
    aliens = []
    for r in range(rows):
        input_row = data[ptr]
        ptr +=1
        maze.append(input_row)
        if 'S' in input_row:
            start = (r, input_row.index('S'))
        for char_index, char in enumerate(input_row):
            if char == 'A':
                aliens.append((r, char_index))
    mazes.append(maze)
    nodes_found = [start] + aliens
    nodes_in_mazes.append(nodes_found)
    number_of_nodes_in_maze.append(len(nodes_found))
    # 预存坐标到ID的映射
    pos_to_id = {pos: idx for idx, pos in enumerate(nodes_found)}
    pos_to_id_list.append(pos_to_id)

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