并查集与线段树

#数据结构 · #并查集 · #Union-Find · #线段树 · #树状数组 · #区间查询

并查集和线段树是解决连通性问题和区间查询问题的两种利器。并查集以近乎 O(1) 的时间处理动态连通性，线段树以 O(log n) 处理区间修改和查询。

一、并查集（Union-Find / Disjoint Set）

核心操作

操作	作用
Find(x)	找到 x 所属集合的代表元素（根）
Union(x, y)	合并 x 和 y 所在的集合
Connected(x, y)	判断 x 和 y 是否在同一集合

基础实现

type UnionFind struct {
    parent []int
}

func NewUnionFind(n int) *UnionFind {
    uf := &UnionFind{parent: make([]int, n)}
    for i := range uf.parent {
        uf.parent[i] = i
    }
    return uf
}

func (uf *UnionFind) Find(x int) int {
    for uf.parent[x] != x {
        x = uf.parent[x]
    }
    return x
}

func (uf *UnionFind) Union(x, y int) {
    rootX, rootY := uf.Find(x), uf.Find(y)
    if rootX != rootY {
        uf.parent[rootX] = rootY
    }
}

优化：路径压缩 + 按秩合并

type UnionFind struct {
    parent []int
    rank   []int  // 树的高度上界
}

func NewUnionFind(n int) *UnionFind {
    uf := &UnionFind{
        parent: make([]int, n),
        rank:   make([]int, n),
    }
    for i := range uf.parent {
        uf.parent[i] = i
        uf.rank[i] = 1
    }
    return uf
}

// 路径压缩：查询过程中把路径上所有节点直接连到根
func (uf *UnionFind) Find(x int) int {
    if uf.parent[x] != x {
        uf.parent[x] = uf.Find(uf.parent[x])  // 递归压缩
    }
    return uf.parent[x]
}

// 按秩合并：矮树挂到高树下
func (uf *UnionFind) Union(x, y int) {
    rootX, rootY := uf.Find(x), uf.Find(y)
    if rootX == rootY {
        return
    }
    if uf.rank[rootX] < uf.rank[rootY] {
        uf.parent[rootX] = rootY
    } else if uf.rank[rootX] > uf.rank[rootY] {
        uf.parent[rootY] = rootX
    } else {
        uf.parent[rootY] = rootX
        uf.rank[rootX]++
    }
}

复杂度：同时使用两种优化后，Find/Union 的均摊时间复杂度为 O(α(n))，其中 α(n) 是反阿克曼函数，实际中 α(n) ≤ 4，可视为常数。

带权并查集

// 维护节点到父节点的权重（如距离、倍数等）
type WeightedUnionFind struct {
    parent []int
    weight []int  // weight[x] = x 到 parent[x] 的权值
}

func (uf *WeightedUnionFind) Find(x int) int {
    if uf.parent[x] != x {
        root := uf.Find(uf.parent[x])
        uf.weight[x] += uf.weight[uf.parent[x]]  // 路径压缩时累加权值
        uf.parent[x] = root
    }
    return uf.parent[x]
}

func (uf *WeightedUnionFind) Union(x, y, w int) {
    // 将 x 所在集合合并到 y，且 x 到 y 的权值为 w
    rootX, rootY := uf.Find(x), uf.Find(y)
    if rootX == rootY {
        return
    }
    uf.parent[rootX] = rootY
    uf.weight[rootX] = uf.weight[y] - uf.weight[x] + w
}

应用

场景	说明
连通性判断	图中有多少连通分量？两点是否连通？
Kruskal 最小生成树	判断加入边是否成环
朋友圈/岛屿问题	合并连通区域
冗余连接	找到删除后仍连通的边
等式方程的可满足性	变量相等关系 → 并查集

二、线段树（Segment Tree）

结构

线段树是一棵平衡二叉树，每个节点存储一个区间的聚合信息（和、最大值、最小值等）。

数组:     [5, 2, 6, 3, 1, 4, 8, 7]
                      [0-7]: 36
                    /            \
            [0-3]: 16            [4-7]: 20
           /        \           /        \
     [0-1]: 7   [2-3]: 9   [4-5]: 5   [6-7]: 15
      /    \     /    \     /    \     /    \
   [0]5 [1]2 [2]6 [3]3 [4]1 [5]4 [6]8 [7]7

完整实现

type SegmentTree struct {
    tree []int
    n    int
}

func NewSegmentTree(arr []int) *SegmentTree {
    n := len(arr)
    st := &SegmentTree{
        tree: make([]int, 4*n),  // 4 倍空间保证足够
        n:    n,
    }
    st.build(arr, 0, 0, n-1)
    return st
}

func (st *SegmentTree) build(arr []int, node, l, r int) {
    if l == r {
        st.tree[node] = arr[l]
        return
    }
    mid := (l + r) / 2
    st.build(arr, node*2+1, l, mid)
    st.build(arr, node*2+2, mid+1, r)
    st.tree[node] = st.tree[node*2+1] + st.tree[node*2+2]
}

// 单点更新：arr[idx] = val
func (st *SegmentTree) Update(idx, val int) {
    st.update(0, 0, st.n-1, idx, val)
}

func (st *SegmentTree) update(node, l, r, idx, val int) {
    if l == r {
        st.tree[node] = val
        return
    }
    mid := (l + r) / 2
    if idx <= mid {
        st.update(node*2+1, l, mid, idx, val)
    } else {
        st.update(node*2+2, mid+1, r, idx, val)
    }
    st.tree[node] = st.tree[node*2+1] + st.tree[node*2+2]
}

// 区间查询：[ql, qr] 的和
func (st *SegmentTree) Query(ql, qr int) int {
    return st.query(0, 0, st.n-1, ql, qr)
}

func (st *SegmentTree) query(node, l, r, ql, qr int) int {
    if ql <= l && r <= qr {
        return st.tree[node]  // 完全覆盖
    }
    if r < ql || l > qr {
        return 0  // 无交集
    }
    mid := (l + r) / 2
    left := st.query(node*2+1, l, mid, ql, qr)
    right := st.query(node*2+2, mid+1, r, ql, qr)
    return left + right
}

懒惰标记（区间更新）

type LazySegmentTree struct {
    tree []int
    lazy []int  // 懒惰标记
    n    int
}

// 区间加值
func (st *LazySegmentTree) RangeAdd(ql, qr, val int) {
    st.rangeAdd(0, 0, st.n-1, ql, qr, val)
}

func (st *LazySegmentTree) pushDown(node, l, r int) {
    if st.lazy[node] != 0 {
        mid := (l + r) / 2
        // 下推给左孩子
        st.tree[node*2+1] += st.lazy[node] * (mid - l + 1)
        st.lazy[node*2+1] += st.lazy[node]
        // 下推给右孩子
        st.tree[node*2+2] += st.lazy[node] * (r - mid)
        st.lazy[node*2+2] += st.lazy[node]
        // 清除当前标记
        st.lazy[node] = 0
    }
}

func (st *LazySegmentTree) rangeAdd(node, l, r, ql, qr, val int) {
    if ql <= l && r <= qr {
        st.tree[node] += val * (r - l + 1)
        st.lazy[node] += val
        return
    }
    st.pushDown(node, l, r)
    mid := (l + r) / 2
    if ql <= mid {
        st.rangeAdd(node*2+1, l, mid, ql, qr, val)
    }
    if qr > mid {
        st.rangeAdd(node*2+2, mid+1, r, ql, qr, val)
    }
    st.tree[node] = st.tree[node*2+1] + st.tree[node*2+2]
}

三、树状数组（Fenwick Tree / BIT）

与线段树对比

维度	线段树	树状数组
单点更新	O(log n)	O(log n)
区间查询（前缀和）	O(log n)	O(log n)
区间更新	O(log n) + lazy	需差分技巧
空间	4n	n+1
代码量	多	少 ✅
适用范围	任意可合并的运算	仅可逆运算（和、异或）

实现

type FenwickTree struct {
    tree []int
    n    int
}

func NewFenwickTree(n int) *FenwickTree {
    return &FenwickTree{tree: make([]int, n+1), n: n}
}

// 单点加值：arr[idx] += delta
func (ft *FenwickTree) Add(idx, delta int) {
    for i := idx + 1; i <= ft.n; i += i & -i {
        ft.tree[i] += delta
    }
}

// 前缀和查询：[0..idx]
func (ft *FenwickTree) Sum(idx int) int {
    sum := 0
    for i := idx + 1; i > 0; i -= i & -i {
        sum += ft.tree[i]
    }
    return sum
}

// 区间和：[l..r]
func (ft *FenwickTree) RangeSum(l, r int) int {
    return ft.Sum(r) - ft.Sum(l-1)
}

关键技巧：i & -i 获取最低位的 1，这是树状数组的核心：

i = 6 (0b0110) → i & -i = 2 (0b0010)
i = 7 (0b0111) → i & -i = 1 (0b0001)

应用

问题	方案
区间和查询	树状数组基础用法
求逆序对	离散化 + 树状数组
区间更新 + 单点查询	差分数组 + 树状数组
区间更新 + 区间查询	双树状数组（维护 Δ 和 i·Δ）
二维偏序问题	二维树状数组

树状数组的区间更新 + 区间查询

标准树状数组只能做"单点更新 + 前缀和查询"。要实现"区间更新 + 区间查询"，需要差分 + 双树状数组的技巧：

// 双树状数组：支持区间加法和区间求和
type RangeFenwick struct {
    tree1, tree2 []int64 // tree1 维护 d[i], tree2 维护 i*d[i]
    n            int
}

func (rf *RangeFenwick) add(idx int, val int64) {
    for i := idx + 1; i <= rf.n; i += i & -i {
        rf.tree1[i] += val
        rf.tree2[i] += val * int64(idx)
    }
}

// 区间 [l, r] 加 v
func (rf *RangeFenwick) RangeAdd(l, r int, v int64) {
    rf.add(l, v)
    rf.add(r+1, -v) // 差分：右端点之后减去
}

// 前缀和 sum[0..idx] = (idx+1) * Σd[i] - Σ(i*d[i])
func (rf *RangeFenwick) prefixSum(idx int) int64 {
    var sum1, sum2 int64
    for i := idx + 1; i > 0; i -= i & -i {
        sum1 += rf.tree1[i]
        sum2 += rf.tree2[i]
    }
    return int64(idx+1)*sum1 - sum2
}

func (rf *RangeFenwick) RangeSum(l, r int) int64 {
    return rf.prefixSum(r) - rf.prefixSum(l-1)
}

数学推导：

区间 [l, r] 加上 v → 对前缀和的影响:
  对于 i < l:   无影响
  对于 l ≤ i ≤ r: sum[i] += (i-l+1)*v = v*i + v*(1-l)
  对于 i > r:    sum[i] += (r-l+1)*v = 常数

看出 sum[i] 可以表示为 a*i + b 的形式
→ 维护两个 BIT: BIṯ 存放系数 a, BIṮ 存放常数项 b → O(log n) 实现区间更新+查询

线段树高级应用

懒传播（Lazy Propagation）详解

当需要区间更新时（如给区间 [l, r] 所有元素加 x），不能一个个更新 O(n log n)——需要用懒标记"推迟"更新。

// 懒传播线段树：区间加法 + 区间求和
type SegTreeLazy struct {
    tree, lazy []int
    n          int
}

func (st *SegTreeLazy) pushDown(node, l, r int) {
    if st.lazy[node] == 0 {
        return
    }
    mid := (l + r) / 2
    val := st.lazy[node]

    // 更新左右子节点的 tree 值
    st.tree[node*2] += val * (mid - l + 1)
    st.tree[node*2+1] += val * (r - mid)

    // 下传懒标记
    st.lazy[node*2] += val
    st.lazy[node*2+1] += val

    // 清除当前节点的懒标记
    st.lazy[node] = 0
}

// 区间 [ql, qr] 加 val
func (st *SegTreeLazy) update(node, l, r, ql, qr, val int) {
    if ql <= l && r <= qr {
        st.tree[node] += val * (r - l + 1) // 整个区间被覆盖
        st.lazy[node] += val               // 打懒标记，不再下传
        return
    }
    st.pushDown(node, l, r) // 需要下钻，先传懒标记
    mid := (l + r) / 2
    if ql <= mid {
        st.update(node*2, l, mid, ql, qr, val)
    }
    if qr > mid {
        st.update(node*2+1, mid+1, r, ql, qr, val)
    }
    st.tree[node] = st.tree[node*2] + st.tree[node*2+1] // 向上合并
}

懒传播的核心思想：

普通的区间更新 [2, 5] + x:
  → 需要更新叶子节点 [2], [3], [4], [5] → O(log n) 每个 → O(n log n) ❌

懒传播:
  [2, 5] 覆盖了节点 [2,3] 和 [4,5] 两个区间 → 直接在这两个节点打标记
  → 只有当后续查询需要深入时才下传标记 → O(log n) ✅

关键: 标记只下传到"需要访问的路径"上，其余节点保持"脏"状态

可持久化线段树（主席树）

可持久化线段树保留每次修改的历史版本。核心技巧：只复制修改路径上的节点，其余节点共享。

mermaid

flowchart TD
    subgraph V1["版本 1: 修改位置 2"]
        R1["root1 → [1,4]=10"]
        R1 --> L1["[1,2]=6"]
        R1 --> R1R["[3,4]=4 (共享)"]
        L1 --> L1L["[1,1]=2 (共享)"]
        L1 --> L1R["[2,2]=4 (复制,新值)"]
    end

    subgraph V0["版本 0: 初始"]
        R0["root0 → [1,4]=8"]
        R0 --> L0["[1,2]=4"]
        R0 --> R0R["[3,4]=4"]
        L0 --> L0L["[1,1]=2"]
        L0 --> L0R["[2,2]=2"]
    end

    V0 -.->|"共享"| V1

// 可持久化线段树节点
type Node struct {
    left, right *Node // 左右子节点指针（不是数组索引）
    sum         int
}

// 更新返回新版本的根节点（旧版本不变）
func update(prev *Node, l, r, idx, val int) *Node {
    cur := &Node{sum: prev.sum + val}
    if l == r {
        return cur
    }
    mid := (l + r) / 2
    if idx <= mid {
        cur.left = update(prev.left, l, mid, idx, val)
        cur.right = prev.right // 右子树共享旧版本
    } else {
        cur.left = prev.left   // 左子树共享旧版本
        cur.right = update(prev.right, mid+1, r, idx, val)
    }
    return cur
}
// 每次更新创建 O(log n) 个新节点，其余共享 → 空间 O(n + q log n)

应用：

问题	方法	复杂度
区间第 K 小	主席树：第 i 个版本 = 前 i 个元素的权值线段树	O(log n)
树上路径第 K 小	主席树 + LCA	O(log n)
历史版本查询	保留每次修改的 root	O(log n)


---

## 带权并查集（Weighted Union-Find）

> 当不仅需要知道连通性，还需要维护节点之间的**关系权重**（距离、比值、差值）时，使用带权并查集。

```go
// WeightedUnionFind: 维护每个节点到根节点的"权重"
// 典型应用: 变量除法 (a/b=c, b/d=e → a/d=?), 食物链
type WeightedUnionFind struct {
    parent []int
    weight []float64  // weight[x] = x / parent[x] (x 到父节点的比值)
}

func NewWeightedUnionFind(n int) *WeightedUnionFind {
    uf := &WeightedUnionFind{
        parent: make([]int, n),
        weight: make([]float64, n),
    }
    for i := range uf.parent {
        uf.parent[i] = i
        uf.weight[i] = 1.0  // 自己到自己的比值为 1
    }
    return uf
}

// Find: 返回根节点 + 当前节点到根的权重（路径压缩 + 权重更新）
func (uf *WeightedUnionFind) Find(x int) (int, float64) {
    if uf.parent[x] != x {
        root, w := uf.Find(uf.parent[x])
        uf.parent[x] = root
        uf.weight[x] *= w  // 链式法则: weight[x→root] = weight[x→parent] × weight[parent→root]
    }
    return uf.parent[x], uf.weight[x]
}

// Union: 合并 x 和 y，满足 x / y = value
func (uf *WeightedUnionFind) Union(x, y int, value float64) bool {
    rootX, wX := uf.Find(x) // wX = x / rootX
    rootY, wY := uf.Find(y) // wY = y / rootY

    if rootX == rootY {
        // 已连通 → 验证一致性: x/y 应该等于 wX / wY
        return math.Abs(wX/wY-value) < 1e-9
    }

    // 挂 rootX 到 rootY:
    // 需要满足: x / y = value
    // 即: (wX × rootX) / (wY × rootY) = value
    // → rootX / rootY = value × wY / wX
    uf.parent[rootX] = rootY
    uf.weight[rootX] = value * wY / wX
    return true
}

// Query: 查询 x / y 的值（如果连通），否则返回错误
func (uf *WeightedUnionFind) Query(x, y int) (float64, bool) {
    rootX, wX := uf.Find(x)
    rootY, wY := uf.Find(y)
    if rootX != rootY {
        return 0, false
    }
    return wX / wY, true // x/y = (x/root) / (y/root)
}

// 示例: LeetCode 399 — 除法求值
func calcEquation(equations [][]string, values []float64, queries [][]string) []float64 {
    uf := NewWeightedUnionFind(40) // 最多 20 对变量, 40 个不同变量
    idx := make(map[string]int)
    id := 0

    for i, eq := range equations {
        a, b := eq[0], eq[1]
        if _, ok := idx[a]; !ok { idx[a] = id; id++ }
        if _, ok := idx[b]; !ok { idx[b] = id; id++ }
        uf.Union(idx[a], idx[b], values[i]) // a / b = values[i]
    }

    res := make([]float64, len(queries))
    for i, q := range queries {
        ia, okA := idx[q[0]]
        ib, okB := idx[q[1]]
        if !okA || !okB {
            res[i] = -1.0
        } else if v, ok := uf.Query(ia, ib); ok {
            res[i] = v
        } else {
            res[i] = -1.0
        }
    }
    return res
}

mermaid

flowchart TB
    subgraph Example["带权并查集示例: a/b=2, b/c=3"]
        direction TB
        Step1["初始: a, b, c 各自独立<br/>weight[i]=1.0"]
        Step2["Union(a, b, 2.0) → a/b=2<br/>parent[a]=b, weight[a]=2.0"]
        Step3["Union(b, c, 3.0) → b/c=3<br/>parent[b]=c, weight[b]=3.0"]
        Step4["Query(a, c):<br/>Find(a) → a→b→c<br/>weight[a] = 2.0 × 3.0 = 6.0<br/>weight[c] = 1.0<br/>→ a/c = 6.0 ✅"]
    end

    Step1 --> Step2 --> Step3 --> Step4

并查集变种	额外信息	典型应用
基础并查集	只维护连通性	朋友圈、岛屿数量
带权并查集	节点间的数值关系	变量除法、方程求解
种类并查集	节点间的类别关系	食物链、二分图判定

问题涉及连通性 / 集合合并？
  → 并查集

问题涉及区间查询 / 区间更新？
  ├── 操作可逆（求和、异或）且只做单点更新？
  │     → 树状数组（最简单）
  ├── 需要区间更新、区间查询、维护最值？
  │     → 线段树
  └── 操作不可逆（最大值、最小值）
        → 线段树

参考

[CLRS 算法导论第21章用于不相交集合的数据结构]
线段树详解
树状数组详解

并查集与线段树 ​

一、并查集（Union-Find / Disjoint Set） ​

核心操作 ​

基础实现 ​

优化：路径压缩 + 按秩合并 ​

带权并查集 ​

应用 ​

二、线段树（Segment Tree） ​

结构 ​

完整实现 ​

懒惰标记（区间更新） ​

三、树状数组（Fenwick Tree / BIT） ​

与线段树对比 ​

实现 ​