Medium

题目描述

给定一个整数 n 和一个有 n 个节点的无向树,节点编号从 0 到 n - 1。树由长度为 n - 1 的二维数组 edges 表示,其中 edges[i] = [ui, vi] 表示 ui 和 vi 之间有一条无向边。

还给定三个不同的目标节点 x、y 和 z。

对于树中的任何节点 u:

  • 设 dx 为从 u 到节点 x 的距离
  • 设 dy 为从 u 到节点 y 的距离
  • 设 dz 为从 u 到节点 z 的距离

如果这三个距离形成勾股三元组,则称节点 u 为特殊节点。

返回树中特殊节点的数量。

勾股三元组由三个整数 a、b、c 组成,当按升序排列时满足 a² + b² = c²。

树中两个节点之间的距离是连接它们的唯一路径上的边数。

示例 1:

输入:n = 4, edges = [[0,1],[0,2],[0,3]], x = 1, y = 2, z = 3
输出:3
解释:
对于每个节点,计算它到节点 x = 1、y = 2、z = 3 的距离。
- 节点 0 的距离为 1、1、1。排序后为 1、1、1,不满足勾股定理。
- 节点 1 的距离为 0、2、2。排序后为 0、2、2。因为 0² + 2² = 2²,节点 1 是特殊的。
- 节点 2 的距离为 2、0、2。排序后为 0、2、2。因为 0² + 2² = 2²,节点 2 是特殊的。
- 节点 3 的距离为 2、2、0。排序后为 0、2、2。也满足勾股定理。
因此节点 1、2、3 是特殊的,答案是 3。

示例 2:

输入:n = 4, edges = [[0,1],[1,2],[2,3]], x = 0, y = 3, z = 2
输出:0

示例 3:

输入:n = 4, edges = [[0,1],[1,2],[1,3]], x = 1, y = 3, z = 0
输出:1

约束:

  • 4 <= n <= 10⁵
  • edges.length == n - 1
  • edges[i] = [ui, vi]
  • 0 <= ui, vi, x, y, z <= n - 1
  • x、y 和 z 两两不同
  • 输入保证 edges 表示一个有效的树

解题思路

解题思路

这道题需要我们找到树中满足勾股定理的特殊节点。解题的关键是计算每个节点到三个目标节点 x、y、z 的距离。

核心思路:

  1. 距离计算:使用 BFS(广度优先搜索)分别从节点 x、y、z 出发,计算它们到所有其他节点的距离
  2. 构建邻接表:将边数组转换为邻接表表示,便于 BFS 遍历
  3. 勾股定理验证:对每个节点,获取它到三个目标节点的距离,排序后检查是否满足 a² + b² = c²

算法步骤:

  1. 根据 edges 构建图的邻接表
  2. 分别从 x、y、z 三个节点进行 BFS,得到三个距离数组
  3. 遍历所有节点,对每个节点:
    • 获取到 x、y、z 的距离
    • 将距离排序为 a ≤ b ≤ c
    • 检查是否满足 a² + b² = c²
  4. 统计满足条件的特殊节点数量

BFS 实现要点:

  • 使用队列进行层次遍历
  • 用 visited 数组避免重复访问
  • 距离从 0 开始逐层递增

时间复杂度主要由三次 BFS 决定,每次 BFS 访问所有节点和边,因此整体复杂度为 O(n)。

代码实现

class Solution {
public:
    vector<int> bfs(const vector<vector<int>>& graph, int start, int n) {
        vector<int> dist(n, -1);
        queue<int> q;
        q.push(start);
        dist[start] = 0;
        
        while (!q.empty()) {
            int node = q.front();
            q.pop();
            
            for (int neighbor : graph[node]) {
                if (dist[neighbor] == -1) {
                    dist[neighbor] = dist[node] + 1;
                    q.push(neighbor);
                }
            }
        }
        
        return dist;
    }
    
    int specialNodes(int n, vector<vector<int>>& edges, int x, int y, int z) {
        // Build adjacency list
        vector<vector<int>> graph(n);
        for (const auto& edge : edges) {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }
        
        // Calculate distances from x, y, z to all nodes
        vector<int> distX = bfs(graph, x, n);
        vector<int> distY = bfs(graph, y, n);
        vector<int> distZ = bfs(graph, z, n);
        
        int count = 0;
        for (int i = 0; i < n; i++) {
            vector<int> distances = {distX[i], distY[i], distZ[i]};
            sort(distances.begin(), distances.end());
            
            int a = distances[0], b = distances[1], c = distances[2];
            if (a * a + b * b == c * c) {
                count++;
            }
        }
        
        return count;
    }
};
class Solution:
    def specialNodes(self, n: int, edges: List[List[int]], x: int, y: int, z: int) -> int:
        from collections import deque, defaultdict
        
        def bfs(graph, start):
            dist = [-1] * n
            queue = deque([start])
            dist[start] = 0
            
            while queue:
                node = queue.popleft()
                for neighbor in graph[node]:
                    if dist[neighbor] == -1:
                        dist[neighbor] = dist[node] + 1
                        queue.append(neighbor)
            
            return dist
        
        # Build adjacency list
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        
        # Calculate distances from x, y, z to all nodes
        dist_x = bfs(graph, x)
        dist_y = bfs(graph, y)
        dist_z = bfs(graph, z)
        
        count = 0
        for i in range(n):
            distances = sorted([dist_x[i], dist_y[i], dist_z[i]])
            a, b, c = distances
            if a * a + b * b == c * c:
                count += 1
        
        return count
public class Solution {
    public int SpecialNodes(int n, int[][] edges, int x, int y, int z) {
        // Build adjacency list
        List<int>[] graph = new List<int>[n];
        for (int i = 0; i < n; i++) {
            graph[i] = new List<int>();
        }
        
        foreach (var edge in edges) {
            graph[edge[0]].Add(edge[1]);
            graph[edge[1]].Add(edge[0]);
        }
        
        // Calculate distances from x, y, z to all nodes
        int[] distX = BFS(graph, x, n);
        int[] distY = BFS(graph, y, n);
        int[] distZ = BFS(graph, z, n);
        
        int count = 0;
        for (int i = 0; i < n; i++) {
            int[] distances = {distX[i], distY[i], distZ[i]};
            Array.Sort(distances);
            
            int a = distances[0], b = distances[1], c = distances[2];
            if (a * a + b * b == c * c) {
                count++;
            }
        }
        
        return count;
    }
    
    private int[] BFS(List<int>[] graph, int start, int n) {
        int[] dist = new int[n];
        Array.Fill(dist, -1);
        Queue<int> queue = new Queue<int>();
        queue.Enqueue(start);
        dist[start] = 0;
        
        while (queue.Count > 0) {
            int node = queue.Dequeue();
            foreach (int neighbor in graph[node]) {
                if (dist[neighbor] == -1) {
                    dist[neighbor] = dist[node] + 1;
                    queue.Enqueue(neighbor);
                }
            }
        }
        
        return dist;
    }
}
var specialNodes = function(n, edges, x, y, z) {
    const adj = Array.from({ length: n }, () => []);
    for (const [u, v] of edges) {
        adj[u].push(v);
        adj[v].push(u);
    }
    
    const bfs = (start) => {
        const dist = new Array(n).fill(-1);
        const queue = [start];
        dist[start] = 0;
        let head = 0;
        
        while (head < queue.length) {
            const node = queue[head++];
            for (const neighbor of adj[node]) {
                if (dist[neighbor] === -1) {
                    dist[neighbor] = dist[node] + 1;
                    queue.push(neighbor);
                }
            }
        }
        return dist;
    };
    
    const distX = bfs(x);
    const distY = bfs(y);
    const distZ = bfs(z);
    
    let count = 0;
    for (let i = 0; i < n; i++) {
        const distances = [distX[i], distY[i], distZ[i]].sort((a, b) => a - b);
        const [a, b, c] = distances;
        if (a * a + b * b === c * c) {
            count++;
        }
    }
    
    return count;
};

复杂度分析

复杂度大小说明
时间复杂度O(n)需要进行三次 BFS,每次 BFS 访问所有 n 个节点和 n-1 条边,总时间复杂度为 O(3n) = O(n)
空间复杂度O(n)需要存储邻接表 O(n)、三个距离数组 O(3n) 和 BFS 队列 O(n),总空间复杂度为 O(n)