網易首頁 > 網易號 > 正文申請入駐

2026-05-04：樹組的交互代價總和。用go語言，給定一個整數 n，以及一棵有 n 個節點的無向樹，節點編號為 0 到 n-1。樹的結構由數組 edges

2026-05-04 00:09:38　來源: moonfdd

北京舉報

分享至

2026-05-04：樹組的交互代價總和。用go語言，給定一個整數 n，以及一棵有 n 個節點的無向樹，節點編號為 0 到 n-1。樹的結構由數組 edges 表示：數組長度為 n-1，其中 edges[i] = [u, v] 表示節點 u 與節點 v 之間有一條無向邊。

另給定一個數組 group，長度為 n。group[i] 表示節點 i 所屬的組。

如果兩個節點 u 和 v 滿足 group[u] == group[v]，則稱它們屬于同一組。由于是樹結構，任意兩個節點之間都存在且僅存在一條唯一路徑。所謂交互代價定義為：這條唯一路徑上包含的邊的數量。

目標：枚舉所有無序且不同的節點對 (u, v)，要求它們同組（group[u] == group[v]）。把這些節點對的交互代價全部累加并返回總和；若不存在滿足條件的節點對，則返回 0。

1 <= n <= 100000。

edges.length == n - 1。

edges[i] = [ui, vi]。

0 <= ui, vi <= n - 1。

group.length == n。

1 <= group[i] <= 20。

輸入保證 edges 表示一棵有效的樹。

輸入： n = 3, edges = [[0,1],[1,2]], group = [3,2,3]。

輸出： 2。

解釋：

節點 0 和節點 2 屬于組 3，它們之間的交互代價為 2。

節點 1 屬于不同的組，因此沒有有效的節點對。

總交互代價為 2。

題目來自力扣3786。

代碼執行全流程詳細拆解第一步：構建原始樹的鄰接表

1. 初始化一個長度為n的二維數組，作為無向樹的鄰接表；
2. 遍歷所有邊，把每條邊的兩個節點互相添加到對方的鄰接列表中；
3. 作用：讓程序可以快速訪問每個節點的所有相鄰節點，為后續樹的遍歷做準備。

第二步：預處理樹的基礎信息（DFS + 倍增LCA）

這一步是為了快速求任意兩點的最近公共祖先（LCA）和兩點間路徑長度，是樹問題的核心預處理。

子步驟1：DFS遍歷樹，記錄核心信息

1. 定義遞歸函數，從根節點0開始遍歷整棵樹；
2. 記錄每個節點的DFS時間戳（dfn）：用于后續給節點排序，是構建虛樹的關鍵；
3. 記錄每個節點的父節點（第一層祖先）；
4. 記錄每個節點的深度（dep）：根節點深度為0，子節點深度=父節點+1；
5. 深度的作用：兩點間路徑邊數 =深度[u] + 深度[v] - 2×深度[LCA(u,v)]。

子步驟2：倍增法預處理祖先數組

1. 計算樹的最大深度對應的二進制位數，確定倍增的層數；
2. 預處理每個節點的2^i級祖先（2級、4級、8級...祖先）；
3. 作用：實現O(logn)時間查詢任意兩個節點的最近公共祖先（LCA）。

子步驟3：封裝LCA查詢函數

1. 封裝兩個工具函數：

? 把節點向上提升到指定深度；
? 求任意兩個節點的最近公共祖先；

2. 這是后續計算路徑長度、構建虛樹的基礎工具。

第三步：按分組歸類所有節點

1. 創建一個哈希表（字典），key=組號，value=該組所有節點的列表；
2. 遍歷所有節點，把每個節點按照group數組的值，放入對應組的列表中；
3. 作用：后續只需要逐組計算，組號最多20個，極大減少計算量。

第四步：對每一組單獨構建「虛樹」（核心優化）

因為組號只有20個，我們逐個組處理，每組獨立計算貢獻：

虛樹作用：只保留當前組的節點 + 這些節點之間路徑的必要公共祖先，剔除無關節點，把大樹壓縮成小樹，大幅降低計算量。

單組構建虛樹的完整步驟

1.按DFS時間戳排序：把當前組的所有節點，按照第一步記錄的dfn從小到大排序；
2.初始化棧和虛樹：用根節點作為棧的初始元素，清空虛樹結構；
3.標記關鍵節點：把當前組的節點標記為「真實關鍵節點」；
4.棧+LCA構建虛樹：

? 遍歷排序后的每個節點；
? 計算棧頂節點與當前節點的LCA（路徑拐點）；
? 不斷回溯棧，給虛樹添加邊，直到棧頂深度小于LCA深度；
? 如果LCA不在棧中，將其加入棧和虛樹；
? 最后把當前節點入棧；

5.收尾加邊：遍歷結束后，把棧中剩余節點依次連邊，完成虛樹構建。

第五步：在虛樹上DFS計算本組的交互代價（貢獻法）

這是計算答案的核心步驟，使用貢獻法：不枚舉所有節點對（會超時），而是計算每條邊被多少對節點經過，總代價 = 邊數 × 經過的節點對數。

計算步驟

1. 定義遞歸DFS函數，遍歷當前組的虛樹；
2. 遞歸統計每個子樹中當前組的節點數量；
3. 對于虛樹上的每一條邊：

? 邊的實際長度 = 子節點深度 - 父節點深度（對應原始樹的邊數）；
? 設子樹內有sz個本組節點，本組總節點數為total；
? 這條邊會被sz × (total - sz)對節點經過；
? 本組總代價 += 邊長度 ×sz × (total - sz)；

4. 把本組的代價累加到全局答案中；

5. 一組計算完成后，重置虛樹，開始處理下一組。

第六步：所有組計算完成，返回最終答案

1. 遍歷完所有組（最多20組）；
2. 全局累加的結果就是所有同組節點對的交互代價總和；
3. 示例中僅組3貢獻了2，最終輸出2。

時間復雜度 & 額外空間復雜度分析一、總時間復雜度

O(n × logn + G × k × logk)
拆解說明：

1.預處理LCA：DFS遍歷樹是O(n)，倍增預處理是O(n × logn)；
2.分組歸類：O(n)；
3.構建虛樹+計算貢獻：

? 組數量G ≤ 20（題目限定）；
? 每組節點數k，排序O(k logk)，構建虛樹+DFSO(k)；
? 所有組總耗時O(n logn)；

4. 整體主導項：O(n × logn)，完全滿足n=1e5的時間要求。

二、總額外空間復雜度

O(n × logn)
拆解說明：

1. 鄰接表：O(n)；
2. 倍增祖先數組：O(n × 17)（log?(1e5)≈17），是核心空間開銷；
3. DFN、深度數組、虛樹、棧、哈希表：均為O(n)；
4. 整體空間復雜度由倍增數組主導：O(n × logn)。

總結

1. 整體流程：建原始樹 → 預處理LCA → 按組分節點 → 每組建虛樹壓縮 → 貢獻法算代價 → 累加答案；
2. 核心優化：利用group[i]≤20的限定，逐組處理+虛樹壓縮，避免暴力枚舉節點對；
3. 時間復雜度：O(n logn)，高效處理1e5節點；
4. 空間復雜度：O(n logn)，符合算法題常規空間要求。

Go完整代碼如下：

package main

 import (
    "fmt"
    "math/bits"
    "slices"
)

 func interactionCosts(n int, edges [][]int, group []int) (ans int64) {
    g := make([][]int, n)
    for _, e := range edges {
        v, w := e[0], e[1]
        g[v] = append(g[v], w)
        g[w] = append(g[w], v)
    }

     dfn := make([]int, n)
    ts := 0
    pa := make([][17]int, n)
    dep := make([]int, n)
    var build func(int, int)
    build = func(v, p int) {
        dfn[v] = ts
        ts++
        pa[v][0] = p
        for _, w := range g[v] {
            if w != p {
                dep[w] = dep[v] + 1
                build(w, v)
            }
        }
    }
    build(0, -1)
    mx := bits.Len(uint(n))
    for i := range mx - 1 {
        for v := range pa {
            p := pa[v][i]
            if p != -1 {
                pa[v][i+1] = pa[p][i]
            } else {
                pa[v][i+1] = -1
            }
        }
    }
    uptoDep := func(v, d int)int {
        for k := uint32(dep[v] - d); k > 0; k &= k - 1 {
            v = pa[v][bits.TrailingZeros32(k)]
        }
        return v
    }
    getLCA := func(v, w int)int {
        if dep[v] > dep[w] {
            v, w = w, v
        }
        w = uptoDep(w, dep[v])
        if w == v {
            return v
        }
        for i := mx - 1; i >= 0; i-- {
            pv, pw := pa[v][i], pa[w][i]
            if pv != pw {
                v, w = pv, pw
            }
        }
        return pa[v][0]
    }

     nodesMap := map[int][]int{}
    for i, x := range group {
        nodesMap[x] = append(nodesMap[x], i)
    }

     vt := make([][]int, n)   // 虛樹
    isNode := make([]int, n) // 用來區分是關鍵節點還是 LCA
    for i := range isNode {
        isNode[i] = -1
    }
    addVtEdge := func(v, w int) {
        vt[v] = append(vt[v], w) // 往虛樹上添加一條有向邊
    }
    const root = 0
    st := []int{root} // 用根節點作為棧底哨兵

     for val, nodes := range nodesMap {
        // 對于相同點權的這一組關鍵節點 nodes，構建虛樹
        slices.SortFunc(nodes, func(a, b int)int { return dfn[a] - dfn[b] })
        vt[root] = vt[root][:0] // 重置虛樹
        st = st[:1]
        for _, v := range nodes {
            isNode[v] = val
            if v == root {
                continue
            }
            vt[v] = vt[v][:0]
            lca := getLCA(st[len(st)-1], v) // 路徑的拐點（LCA）也加到虛樹中
            // 回溯，加邊
            forlen(st) > 1 && dfn[lca] <= dfn[st[len(st)-2]] {
                addVtEdge(st[len(st)-2], st[len(st)-1])
                st = st[:len(st)-1]
            }
            if lca != st[len(st)-1] { // lca 不在棧中（首次遇到）
                vt[lca] = vt[lca][:0]
                addVtEdge(lca, st[len(st)-1])
                st[len(st)-1] = lca // 加到棧中
            }
            st = append(st, v)
        }
        // 最后的回溯，加邊
        for i := 1; i < len(st); i++ {
            addVtEdge(st[i-1], st[i])
        }

         var dfs func(int)int
        dfs = func(v int) (size int) {
            // 如果 isNode[v] != t，那么 v 只是關鍵節點之間路徑上的「拐點」
            if isNode[v] == val {
                size = 1
            }
            for _, w := range vt[v] {
                sz := dfs(w)
                wt := dep[w] - dep[v] // 虛樹邊權
                // 貢獻法
                ans += int64(wt) * int64(sz) * int64(len(nodes)-sz)
                size += sz
            }
            return
        }

         rt := root
        if isNode[rt] != val && len(vt[rt]) == 1 {
            // 注意 root 只是一個哨兵，不一定在虛樹上，得從真正的根節點開始
            rt = vt[rt][0]
        }
        dfs(rt)
    }

     return
}

 func main() {
    n := 3
    edges := [][]int{{0, 1}, {1, 2}}
    group := []int{3, 2, 3}
    result := interactionCosts(n, edges, group)
    fmt.Println(result)
}

Python完整代碼如下：

# -*-coding:utf-8-*-

 import sys
sys.setrecursionlimit(10**6)

 def interactionCosts(n, edges, group):
    ans = 0

     g = [[] for _ in range(n)]
    for v, w in edges:
        g[v].append(w)
        g[w].append(v)

     dfn = [0] * n
    ts = 0
    pa = [[-1] * 17for _ in range(n)]
    dep = [0] * n

     def build(v, p):
        nonlocal ts
        dfn[v] = ts
        ts += 1
        pa[v][0] = p
        for w in g[v]:
            if w != p:
                dep[w] = dep[v] + 1
                build(w, v)

     build(0, -1)

     mx = n.bit_length()
    for i in range(mx - 1):
        for v in range(n):
            p = pa[v][i]
            if p != -1:
                pa[v][i+1] = pa[p][i]
            else:
                pa[v][i+1] = -1

     def uptoDep(v, d):
        k = dep[v] - d
        while k:
            step = (k & -k).bit_length() - 1
            v = pa[v][step]
            k -= (1 << step)
        return v

     def getLCA(v, w):
        if dep[v] > dep[w]:
            v, w = w, v
        w = uptoDep(w, dep[v])
        if w == v:
            return v
        for i in range(mx-1, -1, -1):
            pv, pw = pa[v][i], pa[w][i]
            if pv != pw:
                v, w = pv, pw
        return pa[v][0]

     nodesMap = {}
    for i, x in enumerate(group):
        nodesMap.setdefault(x, []).append(i)

     vt = [[] for _ in range(n)]
    isNode = [-1] * n

     def addVtEdge(v, w):
        vt[v].append(w)

     root = 0
    st = [root]

     for val, nodes in nodesMap.items():
        nodes.sort(key=lambda x: dfn[x])
        vt[root] = []
        st = [root]
        for v in nodes:
            isNode[v] = val
            if v == root:
                continue
            vt[v] = []
            lca = getLCA(st[-1], v)
            while len(st) > 1 and dfn[lca] <= dfn[st[-2]]:
                addVtEdge(st[-2], st[-1])
                st.pop()
            if lca != st[-1]:
                vt[lca] = []
                addVtEdge(lca, st[-1])
                st[-1] = lca
            st.append(v)
        for i in range(1, len(st)):
            addVtEdge(st[i-1], st[i])

         sys.setrecursionlimit(10**6)
        def dfs(v):
            nonlocal ans
            size = 1if isNode[v] == val else0
            for w in vt[v]:
                sz = dfs(w)
                wt = dep[w] - dep[v]
                ans += wt * sz * (len(nodes) - sz)
                size += sz
            return size

         rt = root
        if isNode[rt] != val and len(vt[rt]) == 1:
            rt = vt[rt][0]
        dfs(rt)

     return ans

 if __name__ == "__main__":
    n = 3
    edges = [[0, 1], [1, 2]]
    group = [3, 2, 3]
    result = interactionCosts(n, edges, group)
    print(result)

C++完整代碼如下：

  

  

   

  

  


 using namespace std;

 class Solution {
public:
    long long interactionCosts(int n, vector 
 
 int 
 >>& edges, vector< 
 int 
 >& group) { 
 
 
         long long ans =  
 0 
 ; 
 
 
 
 
          
 // 構建鄰接表 
 
 
         vector 
 
 int 
 >> g(n); 
 
 
          
 for 
  (auto& e : edges) { 
 
 
              
 int 
  v = e[ 
 0 
 ], w = e[ 
 1 
 ]; 
 
 
             g[v].push_back(w); 
 
 
             g[w].push_back(v); 
 
 
         } 
 
 
 
 
          
 // 預處理 DFS 序、深度和倍增祖先 
 
 
         vector< 
 int 
 > dfn(n,  
 0 
 ); 
 
 
          
 int 
  ts =  
 0 
 ; 
 
 
         vector 
 
 int 
 ,  
 17 
 >> pa(n); 
 
 
         vector< 
 int 
 > dep(n,  
 0 
 ); 
 
 
 
 
         function 
 
 int 
 ,  
 int 
 )> build = [&]( 
 int 
  v,  
 int 
  p) { 
 
 
             dfn[v] = ts++; 
 
 
             pa[v][ 
 0 
 ] = p; 
 
 
              
 for 
  ( 
 int 
  w : g[v]) { 
 
 
                  
 if 
  (w != p) { 
 
 
                     dep[w] = dep[v] +  
 1 
 ; 
 
 
                     build(w, v); 
 
 
                 } 
 
 
             } 
 
 
         }; 
 
 
         build( 
 0 
 ,  
 -1 
 ); 
 
 
 
 
          
 int 
  mx =  
 32 
  - __builtin_clz(n);  
 // bits.Len(uint(n)) 
 
 
          
 for 
  ( 
 int 
  i =  
 0 
 ; i < mx -  
 1 
 ; i++) { 
 
 
              
 for 
  ( 
 int 
  v =  
 0 
 ; v < n; v++) { 
 
 
                  
 int 
  p = pa[v][i]; 
 
 
                  
 if 
  (p !=  
 -1 
 ) { 
 
 
                     pa[v][i +  
 1 
 ] = pa[p][i]; 
 
 
                 }  
 else 
  { 
 
 
                     pa[v][i +  
 1 
 ] =  
 -1 
 ; 
 
 
                 } 
 
 
             } 
 
 
         } 
 
 
 
 
          
 // 跳到指定深度 
 
 
         auto uptoDep = [&]( 
 int 
  v,  
 int 
  d) ->  
 int 
  { 
 
 
              
 int 
  k = dep[v] - d; 
 
 
             while (k >  
 0 
 ) { 
 
 
                  
 int 
  step = __builtin_ctz(k); 
 
 
                 v = pa[v][step]; 
 
 
                 k &= k -  
 1 
 ; 
 
 
             } 
 
 
              
 return 
  v; 
 
 
         }; 
 
 
 
 
          
 // 獲取 LCA 
 
 
         auto getLCA = [&]( 
 int 
  v,  
 int 
  w) ->  
 int 
  { 
 
 
              
 if 
  (dep[v] > dep[w]) { 
 
 
                 swap(v, w); 
 
 
             } 
 
 
             w = uptoDep(w, dep[v]); 
 
 
              
 if 
  (w == v)  
 return 
  v; 
 
 
              
 for 
  ( 
 int 
  i = mx -  
 1 
 ; i >=  
 0 
 ; i--) { 
 
 
                  
 int 
  pv = pa[v][i], pw = pa[w][i]; 
 
 
                  
 if 
  (pv != pw) { 
 
 
                     v = pv; 
 
 
                     w = pw; 
 
 
                 } 
 
 
             } 
 
 
              
 return 
  pa[v][ 
 0 
 ]; 
 
 
         }; 
 
 
 
 
          
 // 按點權分組節點 
 
 
          
 map 
 < 
 int 
 , vector< 
 int 
 >> nodesMap; 
 
 
          
 for 
  ( 
 int 
  i =  
 0 
 ; i < n; i++) { 
 
 
             nodesMap[group[i]].push_back(i); 
 
 
         } 
 
 
 
 
          
 // 虛樹 
 
 
         vector 
 
 int 
 >> vt(n); 
 
 
         vector< 
 int 
 > isNode(n,  
 -1 
 ); 
 
 
 
 
         auto addVtEdge = [&]( 
 int 
  v,  
 int 
  w) { 
 
 
             vt[v].push_back(w); 
 
 
         }; 
 
 
 
 
          
 const 
 int 
  root =  
 0 
 ; 
 
 
         vector< 
 int 
 > st; 
 
 
 
 
          
 // 處理每個點權組 
 
 
          
 for 
  (auto& [val, nodes] : nodesMap) { 
 
 
              
 // 按 DFS 序排序 
 
 
             sort(nodes.begin(), nodes.end(), [&]( 
 int 
  a,  
 int 
  b) { 
 
 
                  
 return 
  dfn[a] < dfn[b]; 
 
 
             }); 
 
 
 
 
              
 // 清空虛樹 
 
 
              
 for 
  ( 
 int 
  v : nodes) { 
 
 
                 vt[v].clear(); 
 
 
             } 
 
 
             vt[root].clear(); 
 
 
 
 
             st.clear(); 
 
 
             st.push_back(root); 
 
 
 
 
              
 // 構建虛樹 
 
 
              
 for 
  ( 
 int 
  v : nodes) { 
 
 
                 isNode[v] = val; 
 
 
                  
 if 
  (v == root)  
 continue 
 ; 
 
 
 
 
                 vt[v].clear(); 
 
 
                  
 int 
  lca = getLCA(st.back(), v); 
 
 
 
 
                  
 // 回溯并加邊 
 
 
                 while (st.size() >  
 1 
  && dfn[lca] <= dfn[st[st.size() -  
 2 
 ]]) { 
 
 
                     addVtEdge(st[st.size() -  
 2 
 ], st.back()); 
 
 
                     st.pop_back(); 
 
 
                 } 
 
 
 
 
                  
 if 
  (lca != st.back()) { 
 
 
                     vt[lca].clear(); 
 
 
                     addVtEdge(lca, st.back()); 
 
 
                     st.back() = lca; 
 
 
                 } 
 
 
 
 
                 st.push_back(v); 
 
 
             } 
 
 
 
 
              
 // 添加剩余邊 
 
 
              
 for 
  ( 
 int 
  i =  
 1 
 ; i < st.size(); i++) { 
 
 
                 addVtEdge(st[i -  
 1 
 ], st[i]); 
 
 
             } 
 
 
 
 
              
 // DFS 遍歷虛樹計算貢獻 
 
 
             function< 
 int 
 ( 
 int 
 )> dfs = [&]( 
 int 
  v) ->  
 int 
  { 
 
 
                  
 int 
  size = (isNode[v] == val) ?  
 1 
  :  
 0 
 ; 
 
 
                  
 for 
  ( 
 int 
  w : vt[v]) { 
 
 
                      
 int 
  sz = dfs(w); 
 
 
                      
 int 
  wt = dep[w] - dep[v]; 
 
 
                     ans +=  
 1 
 LL * wt * sz * (nodes.size() - sz); 
 
 
                     size += sz; 
 
 
                 } 
 
 
                  
 return 
  size; 
 
 
             }; 
 
 
 
 
              
 // 找到真正的根節點 
 
 
              
 int 
  rt = root; 
 
 
              
 if 
  (isNode[rt] != val && vt[rt].size() ==  
 1 
 ) { 
 
 
                 rt = vt[rt][ 
 0 
 ]; 
 
 
             } 
 
 
             dfs(rt); 
 
 
         } 
 
 
 
 
          
 return 
  ans; 
 
 
     } 
 
 
 }; 
 
 
 
 
 int 
  main() { 
 
 
      
 int 
  n =  
 3 
 ; 
 
 
     vector 
 
 int 
 >> edges = {{ 
 0 
 ,  
 1 
 }, { 
 1 
 ,  
 2 
 }}; 
 
 
     vector< 
 int 
 > group = { 
 3 
 ,  
 2 
 ,  
 3 
 }; 
 
 
 
 
     Solution solution; 
 
 
     long long result = solution.interactionCosts(n, edges, group); 
 
 
     cout << result << endl; 
 
 
 
 
      
 return 
 0 
 ; 
 
 
 }

我們相信人工智能為普通人提供了一種“增強工具”，并致力于分享全方位的AI知識。在這里，您可以找到最新的AI科普文章、工具評測、提升效率的秘籍以及行業洞察。歡迎關注“福大大架構師每日一題”，發消息可獲得面試資料，讓AI助力您的未來發展。

特別聲明：以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布，本平臺僅提供信息存儲服務。

Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.