Go实现LeetCode 2493超时求助：代码正确但运行过慢

我是物理专业出身的系统管理员，靠兴趣自学编程。我的Go代码能正确解决LeetCode 2493题，但存在运行超时的问题。

解题思路 #

• 使用**并查集（Union-Find）**算法找出图中所有连通子图；
• 对每个连通子图，以每个顶点为根构建BFS树，获取最大深度，同时验证图是否为二分图。

核心逻辑代码 #

func magnificentSets(n int, edges [][]int) int {
    // 构建图
    g := GraphWithCapacity(n)
    set := DisjointSetWiithCapacity(n)
    for _, e := range edges {
        src := e[0] - 1
        dst := e[1] - 1
        if !g.AddEdge(src, dst) {
            fmt.Printf("无法添加边 {%d, %d}\n.", src, dst)
        }
        set.Union(src, dst)
    }

    connected_subgraphs := make(map[int][]int)
    for n := range g.EachNode() {
        p := set.Find(n)
        connected_subgraphs[p] = append(connected_subgraphs[p], n)
    }
    total_maxdepths := 0

    for _, vlist := range connected_subgraphs {
        if len(vlist) == 1 {
            total_maxdepths += 1
            continue
        }

        maxdepth := 0
        for _, v := range vlist {
            _, depthatv, err := bfs(g, v)
            if err != nil {
                fmt.Println(err.Error())
                return -1
            }

            if depthatv > maxdepth {
                maxdepth = depthatv
            }

        }

        total_maxdepths += maxdepth + 1
    }

    return total_maxdepths
}

代码拆分说明 #

并查集实现 #

type disjoint_set struct {
    parent []int
}

func DisjointSetWiithCapacity(n int) disjoint_set {
    parent := make([]int, n)
    for idx := range parent {
        parent[idx] = idx
    }

    return disjoint_set{parent}
}

func (uf *disjoint_set) Find(node int) int {
    cnode := node

    for uf.parent[cnode] != cnode {
        cnode = uf.parent[cnode]
    }
    uf.parent[node] = cnode
    return cnode
}
func (uf *disjoint_set) Union(node_1, node_2 int) bool {
    p1 := uf.Find(node_1)
    p2 := uf.Find(node_2)
    if p1 == p2 {
        return false
    }

    child := max(p1, p2)
    parent := min(p1, p2)
    uf.parent[child] = parent
    return true
}

图数据结构实现 #

type ErrNodeNotFound struct {
    nodeid int
}

func (e ErrNodeNotFound) Error() string {
    return fmt.Sprintf("图中未找到节点 %d。\n", e.nodeid)
}

type graph struct {
    fadjlist [][]int
}

func GraphWithCapacity(cap int) graph {
    return graph{fadjlist: make([][]int, cap)}
}

图操作方法 #

func (g *graph) AddNode() int {
    g.fadjlist = append(g.fadjlist, make([]int, 0))
    return len(g.fadjlist) - 1
}
func (g *graph) HasNode(nodeid int) bool {
    return nodeid < len(g.fadjlist)
}
func (g *graph) EachNode() iter.Seq[int] {
    return func(yield func(int) bool) {
        for i := 0; i < len(g.fadjlist); i += 1 {
            if !yield(i) {
                return
            }
        }
    }
}
func (g *graph) EachEdge() iter.Seq2[int, int] {
    return func(yield func(int, int) bool) {
        for src, adjlist := range g.fadjlist {
            for _, dst := range adjlist {
                if dst > src {
                    continue
                }

                if !yield(src, dst) {
                    return
                }
            }
        }
    }
}
func (g *graph) AddEdge(from, to int) bool {
    src := min(from, to)
    dst := max(from, to)
    if !g.HasNode(src) || !g.HasNode(dst) {
        return false
    }

    if slices.Contains(g.fadjlist[src], dst) {
        return false
    }

    g.fadjlist[src] = append(g.fadjlist[src], dst)
    g.fadjlist[dst] = append(g.fadjlist[dst], src)

    return true
}
func (g *graph) Neighbours(nodeid int) ([]int, error) {
    if !g.HasNode(nodeid) {
        return []int{}, ErrNodeNotFound{nodeid}
    }

    return g.fadjlist[nodeid], nil
}

元图结构及操作 #

type bfsdata struct {
    Color int
    Depth int
}

type metagraph[T any] struct {
    G        *graph
    metadata map[int]T
}

func (mg *metagraph[T]) Data(v int) (T, error) {
    if !mg.G.HasNode(v) {
        return mg.metadata[v], ErrNodeNotFound{v}
    }

    return mg.metadata[v], nil
}
func (mg *metagraph[T]) SetData(v int, data T) (T, error) {
    if !mg.G.HasNode(v) {
        return mg.metadata[v], ErrNodeNotFound{v}
    }

    mg.metadata[v] = data
    return mg.metadata[v], nil
}

BFS核心逻辑 #

func ProcessChildren(mg metagraph[bfsdata], parent int, checked map[int]bool) bool {
    children, e := mg.G.Neighbours(parent)
    if e != nil {
        panic(e)
    }

    mydata, err := mg.Data(parent)
    if err != nil {
        return false
    }

    for _, child := range children {
        if checked[child] {
            continue
        }

        childdata, e := mg.Data(child)
        if e != nil {
            return false
        } else if childdata.Color == mydata.Color {
            return false
        }

        mg.SetData(child, bfsdata{Color: -1 * mydata.Color, Depth: mydata.Depth + 1})

    }

    return true
}


type ErrNotBipartite struct {
    parent int
}

func (e ErrNotBipartite) Error() string {
    return fmt.Sprintf("节点 %d 与同一分区的节点相连", e.parent)
}

func bfs(g graph, root int) (metagraph[bfsdata], int, error) {
    mg := metagraph[bfsdata]{&g, make(map[int]bfsdata)}
    mg.SetData(root, bfsdata{Color: -1, Depth: 0})

    // 以第一个节点为根启动BFS
    tocheck := []int{root}
    maxdepth := 0
    _, e := mg.G.Neighbours(root)
    if e != nil {
        panic(e)
    }
    checked := make(map[int]bool)

    for len(tocheck) > 0 {
        checkme := tocheck[0]
        tocheck = tocheck[1:]

        if checked[checkme] {
            continue
        }

        if !ProcessChildren(mg, checkme, checked) {
            return mg, -1, ErrNotBipartite{checkme}
        }

        mydata, _ := mg.Data(checkme)
        if maxdepth < mydata.Depth {
            maxdepth = mydata.Depth
        }
        nlist, _ := mg.G.Neighbours(checkme)
        tocheck = append(tocheck, nlist...)
        checked[checkme] = true
    }

    return mg, maxdepth, nil
}

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