Description

給定一棵有 n 個節點的樹,每個節點有值 vals[i]。一條「好路徑」的兩端節點值相同,且路徑上所有節點值都不大於端點值。求好路徑的數量(包括長度為 0 的路徑,即每個節點本身)。

Example:

Input: vals = [1,3,2,1,3], edges = [[0,1],[0,2],[2,3],[2,4]] Output: 6

Intuition#

按節點值從小到大加入,用 Union-Find 合併。每次加入一批相同值的節點時,計算新形成的好路徑數。

  • 好路徑的端點值相同,且中間所有節點值不超過端點值
  • 按值從小到大處理節點,確保已加入的節點值都不超過當前值
  • 每次加入同一個值的節點後,計算有多少個在同一連通分量中

Approaches#

⭐1: Union-Find + 排序#

  • 概念: 按節點值排序,依序將節點及其邊加入 Union-Find。每組相同值的節點中,同一連通分量內的 k 個節點貢獻 C(k,2) 條好路徑。
  • 時間複雜度: O(n * α(n))(近似線性)
  • 空間複雜度: O(n)
class Solution {
    private lateinit var parent: IntArray
    private lateinit var rank: IntArray

    fun numberOfGoodPaths(vals: IntArray, edges: Array<IntArray>): Int {
        val n = vals.size
        parent = IntArray(n) { it }
        rank = IntArray(n)

        val graph = Array(n) { mutableListOf<Int>() }
        for ((u, v) in edges) {
            graph[u].add(v)
            graph[v].add(u)
        }

        // 按值分組
        val valueToNodes = TreeMap<Int, MutableList<Int>>()
        for (i in 0 until n) {
            valueToNodes.getOrPut(vals[i]) { mutableListOf() }.add(i)
        }

        var result = n  // 每個節點本身是一條長度 0 的好路徑

        for ((value, nodes) in valueToNodes) {
            // 將當前值的節點與值 <= value 的鄰居合併
            for (node in nodes) {
                for (neighbor in graph[node]) {
                    if (vals[neighbor] <= value) {
                        union(node, neighbor)
                    }
                }
            }

            // 計算同一連通分量中相同值的節點對數
            val countByRoot = HashMap<Int, Int>()
            for (node in nodes) {
                val root = find(node)
                countByRoot[root] = (countByRoot[root] ?: 0) + 1
            }

            for (count in countByRoot.values) {
                result += count * (count - 1) / 2
            }
        }

        return result
    }

    private fun find(x: Int): Int {
        if (parent[x] != x) parent[x] = find(parent[x])
        return parent[x]
    }

    private fun union(x: Int, y: Int) {
        val px = find(x)
        val py = find(y)
        if (px == py) return
        if (rank[px] < rank[py]) parent[px] = py
        else if (rank[px] > rank[py]) parent[py] = px
        else { parent[py] = px; rank[px]++ }
    }
}

2: 按邊排序的 Union-Find#

  • 概念: 將邊按兩端最大值排序,依序加入。處理到值 v 的邊時,統計值為 v 的節點在同一分量中的數量。
  • 時間複雜度: O(n * α(n))
  • 空間複雜度: O(n)
class Solution {
    private lateinit var parent: IntArray
    private lateinit var rank: IntArray

    fun numberOfGoodPaths(vals: IntArray, edges: Array<IntArray>): Int {
        val n = vals.size
        parent = IntArray(n) { it }
        rank = IntArray(n)

        // 按邊的兩端最大值排序
        edges.sortBy { maxOf(vals[it[0]], vals[it[1]]) }

        // 每個分量中最大值的計數
        val maxCount = IntArray(n) { 1 }  // 每個節點初始為自己分量的最大值

        var result = n

        for ((u, v) in edges) {
            val pu = find(u)
            val pv = find(v)
            if (pu == pv) continue

            val maxU = vals[pu]
            val maxV = vals[pv]

            if (maxU == maxV) {
                result += maxCount[pu] * maxCount[pv]
                union(pu, pv)
                val root = find(pu)
                maxCount[root] = maxCount[pu] + maxCount[pv]
            } else if (maxU > maxV) {
                union(pv, pu)
                val root = find(pu)
                maxCount[root] = maxCount[pu]
            } else {
                union(pu, pv)
                val root = find(pv)
                maxCount[root] = maxCount[pv]
            }
        }

        return result
    }

    private fun find(x: Int): Int {
        if (parent[x] != x) parent[x] = find(parent[x])
        return parent[x]
    }

    private fun union(x: Int, y: Int) {
        val px = find(x)
        val py = find(y)
        if (px == py) return
        if (rank[px] < rank[py]) parent[px] = py
        else if (rank[px] > rank[py]) parent[py] = px
        else { parent[py] = px; rank[px]++ }
    }
}

🔑 Takeaways#

  • Pattern: 離線 Union-Find – 按值排序後逐步合併,利用單調性保證合法性
  • 關鍵技巧: 按值從小到大加入確保路徑上的值不超過端點值;每組相同值的 k 個節點在同一分量貢獻 C(k,2) 條好路徑