rust 解法, 执行用时: 57 ms, 内存消耗: 2.4 MB, 提交时间: 2024-06-04 09:37:36

impl Solution {
    pub fn count_pairs_of_connectable_servers(edges: Vec<Vec<i32>>, signal_speed: i32) -> Vec<i32> {
        let n = edges.len() + 1;
        let mut graph = vec![Vec::new(); n];
        for edge in edges.iter() {
            let x = edge[0] as usize;
            let y = edge[1] as usize;
            let cost = edge[2];
            graph[x].push((y, cost));
            graph[y].push((x, cost));
        }

        fn dfs(graph: &Vec<Vec<(usize, i32)>>, p: usize, root: usize, curr: i32, signal_speed: i32) -> i32 {
            let mut res = 0;
            if curr == 0 {
                res += 1;
            }
            for &(v, cost) in &graph[p] {
                if v != root {
                    res += dfs(graph, v, p, (curr + cost) % signal_speed, signal_speed);
                }
            }
            res
        }
        
        let mut ans = vec![0; n];
        for i in 0..n {
            let mut pre = 0;
            for &(v, cost) in &graph[i] {
                let cnt = dfs(&graph, v, i, cost % signal_speed, signal_speed);
                ans[i] += pre * cnt;
                pre += cnt;
            } 
        }
        ans
    }
}

javascript 解法, 执行用时: 347 ms, 内存消耗: 57.7 MB, 提交时间: 2024-06-04 09:37:03

/**
 * @param {number[][]} edges
 * @param {number} signalSpeed
 * @return {number[]}
 */
var countPairsOfConnectableServers = function(edges, signalSpeed) {
    const n = edges.length + 1;
    const graph = Array.from({ length: n }, () => []);
    for (const [u, v, w] of edges) {
        graph[u].push([v, w]);
        graph[v].push([u, w]);
    }

    const dfs = (p, root, curr) => {
        let res = 0;
        if (curr === 0) {
            res++;
        }
        for (const [v, cost] of graph[p]) {
            if (v !== root) {
                res += dfs(v, p, (curr + cost) % signalSpeed);
            }
        }
        return res;
    };

    const res = Array(n).fill(0);
    for (let i = 0; i < n; i++) {
        let pre = 0;
        for (const [v, cost] of graph[i]) {
            const cnt = dfs(v, i, cost % signalSpeed);
            res[i] += pre * cnt;
            pre += cnt;
        }
    }
    return res;
};

cpp 解法, 执行用时: 308 ms, 内存消耗: 35.3 MB, 提交时间: 2024-06-04 09:35:57

class Solution {
public:
    vector<int> countPairsOfConnectableServers(vector<vector<int>> &edges, int signalSpeed) {
        int n = edges.size() + 1;
        vector<vector<pair<int, int>>> g(n);
        for (auto &e : edges) {
            int x = e[0], y = e[1], wt = e[2];
            g[x].push_back({y, wt});
            g[y].push_back({x, wt});
        }

        function<int(int, int, int)> dfs = [&](int x, int fa, int sum) -> int {
            int cnt = sum % signalSpeed == 0;
            for (auto &[y, wt] : g[x]) {
                if (y != fa) {
                    cnt += dfs(y, x, sum + wt);
                }
            }
            return cnt;
        };

        vector<int> ans(n);
        for (int i = 0; i < n; i++) {
            int sum = 0;
            for (auto &[y, wt] : g[i]) {
                int cnt = dfs(y, i, wt);
                ans[i] += cnt * sum;
                sum += cnt;
            }
        }
        return ans;
    }
};

php 解法, 执行用时: 899 ms, 内存消耗: 21.6 MB, 提交时间: 2024-03-06 17:27:06

class Solution {

    /**
     * @param Integer[][] $edges
     * @param Integer $signalSpeed
     * @return Integer[]
     */
    function countPairsOfConnectableServers($edges, $signalSpeed) {
        $n = count($edges) + 1;
        $g = array_fill(0, $n, []);

        foreach ($edges as $e) {
            $x = $e[0];
            $y = $e[1];
            $wt = $e[2];
            $g[$x][] = [$y, $wt];
            $g[$y][] = [$x, $wt];
        }

        $ans = array_fill(0, $n, 0);
        for ($i = 0; $i < $n; $i++) {
            $sum = 0;
            foreach ($g[$i] as $e ) {
                $cnt = $this->dfs($e[0], $i, $e[1], $g, $signalSpeed);
                $ans[$i] += $cnt * $sum;
                $sum += $cnt;
            }
        }
        return $ans;
    }

    private function dfs($x, $fa, $sum, $g, $signalSpeed) {
        $cnt = $sum % $signalSpeed == 0 ? 1 : 0;
        foreach ($g[$x] as $e) {
            $y = $e[0];
            if ( $y != $fa) {
                $cnt += $this->dfs($y, $x, $sum + $e[1], $g, $signalSpeed);
            }
        }
        return $cnt;
    }
}

java 解法, 执行用时: 253 ms, 内存消耗: 44.8 MB, 提交时间: 2024-03-06 17:21:57

class Solution {
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1;
        List<int[]>[] g = new ArrayList[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int[] e : edges) {
            int x = e[0];
            int y = e[1];
            int wt = e[2];
            g[x].add(new int[]{y, wt});
            g[y].add(new int[]{x, wt});
        }

        int[] ans = new int[n];
        for (int i = 0; i < n; i++) {
            int sum = 0;
            for (int[] e : g[i]) {
                int cnt = dfs(e[0], i, e[1], g, signalSpeed);
                ans[i] += cnt * sum;
                sum += cnt;
            }
        }
        return ans;
    }

    private int dfs(int x, int fa, int sum, List<int[]>[] g, int signalSpeed) {
        int cnt = sum % signalSpeed == 0 ? 1 : 0;
        for (int[] e : g[x]) {
            int y = e[0];
            if (y != fa) {
                cnt += dfs(y, x, sum + e[1], g, signalSpeed);
            }
        }
        return cnt;
    }
}

golang 解法, 执行用时: 90 ms, 内存消耗: 6.8 MB, 提交时间: 2024-03-06 17:09:15

func countPairsOfConnectableServers(edges [][]int, signalSpeed int) []int {
	n := len(edges) + 1
	type edge struct{ to, wt int }
	g := make([][]edge, n)
	for _, e := range edges {
		x, y, wt := e[0], e[1], e[2]
		g[x] = append(g[x], edge{y, wt})
		g[y] = append(g[y], edge{x, wt})
	}

	ans := make([]int, n)
	for i, gi := range g {
		var cnt int
		var dfs func(int, int, int)
		dfs = func(x, fa, sum int) {
			if sum%signalSpeed == 0 {
				cnt++
			}
			for _, e := range g[x] {
				if e.to != fa {
					dfs(e.to, x, sum+e.wt)
				}
			}
			return
		}
		sum := 0
		for _, e := range gi {
			cnt = 0
			dfs(e.to, i, e.wt)
			ans[i] += cnt * sum
			sum += cnt
		}
	}
	return ans
}

python3 解法, 执行用时: 1351 ms, 内存消耗: 18.1 MB, 提交时间: 2024-03-06 17:05:05

# 枚举根 + 乘法原理
class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]:
        n = len(edges) + 1
        g = [[] for _ in range(n)]
        for x, y, wt in edges:
            g[x].append((y, wt))
            g[y].append((x, wt))

        def dfs(x: int, fa: int, s: int) -> int:
            cnt = 0 if s % signalSpeed else 1
            for y, wt in g[x]:
                if y != fa:
                    cnt += dfs(y, x, s + wt)
            return cnt

        ans = [0] * n
        for i, gi in enumerate(g):
            s = 0
            for y, wt in gi:
                cnt = dfs(y, i, wt)
                ans[i] += cnt * s
                s += cnt
        return ans