<p>You are given an unrooted weighted tree with <code>n</code> vertices representing servers numbered from <code>0</code> to <code>n - 1</code>, an array <code>edges</code> where <code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>, weight<sub>i</sub>]</code> represents a bidirectional edge between vertices <code>a<sub>i</sub></code> and <code>b<sub>i</sub></code> of weight <code>weight<sub>i</sub></code>. You are also given an integer <code>signalSpeed</code>.</p>

<p>Two servers <code>a</code> and <code>b</code> are <strong>connectable</strong> through a server <code>c</code> if:</p>

<ul>
	<li><code>a &lt; b</code>, <code>a != c</code> and <code>b != c</code>.</li>
	<li>The distance from <code>c</code> to <code>a</code> is divisible by <code>signalSpeed</code>.</li>
	<li>The distance from <code>c</code> to <code>b</code> is divisible by <code>signalSpeed</code>.</li>
	<li>The path from <code>c</code> to <code>b</code> and the path from <code>c</code> to <code>a</code> do not share any edges.</li>
</ul>

<p>Return <em>an integer array</em> <code>count</code> <em>of length</em> <code>n</code> <em>where</em> <code>count[i]</code> <em>is the <strong>number</strong> of server pairs that are <strong>connectable</strong> through</em> <em>the server</em> <code>i</code>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2024/01/21/example22.png" style="width: 438px; height: 243px; padding: 10px; background: #fff; border-radius: .5rem;" />
<pre>
<strong>Input:</strong> edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1
<strong>Output:</strong> [0,4,6,6,4,0]
<strong>Explanation:</strong> Since signalSpeed is 1, count[c] is equal to the number of pairs of paths that start at c and do not share any edges.
In the case of the given path graph, count[c] is equal to the number of servers to the left of c multiplied by the servers to the right of c.
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2024/01/21/example11.png" style="width: 495px; height: 484px; padding: 10px; background: #fff; border-radius: .5rem;" />
<pre>
<strong>Input:</strong> edges = [[0,6,3],[6,5,3],[0,3,1],[3,2,7],[3,1,6],[3,4,2]], signalSpeed = 3
<strong>Output:</strong> [2,0,0,0,0,0,2]
<strong>Explanation:</strong> Through server 0, there are 2 pairs of connectable servers: (4, 5) and (4, 6).
Through server 6, there are 2 pairs of connectable servers: (4, 5) and (0, 5).
It can be shown that no two servers are connectable through servers other than 0 and 6.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>2 &lt;= n &lt;= 1000</code></li>
	<li><code>edges.length == n - 1</code></li>
	<li><code>edges[i].length == 3</code></li>
	<li><code>0 &lt;= a<sub>i</sub>, b<sub>i</sub> &lt; n</code></li>
	<li><code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>, weight<sub>i</sub>]</code><!-- notionvc: a2623897-1bb1-4c07-84b6-917ffdcd83ec --></li>
	<li><code>1 &lt;= weight<sub>i</sub> &lt;= 10<sup>6</sup></code></li>
	<li><code>1 &lt;= signalSpeed &lt;= 10<sup>6</sup></code></li>
	<li>The input is generated such that <code>edges</code> represents a valid tree.</li>
</ul>

class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]:
        '''
        Clarify:
        - return an array of the count of connctable pairs through server i
        - two servers are connectable if the DISTANCE between them is divisible by signalSpeed
        - signalSpeed and weights are positive
        - the input is generated such that you have a valid tree (no strongly connected components)

        Brute Force: Find # of server pairs for each node

        Base cases:
        - if the starting node has an outdegree of 1, there can't be any pairs so return 0


          1   3   5  
        0   2   4 


        For each node run a dfs on two all other pairings of nodes
        TC: O(n + n^3)
        SC: O(n)

        Approach 2:

          1   3   5  
        0   2   4

        Trim leaf nodes (nodes that have an outdegree of 1), then add its parent to the queue
        But keep track of the sums we have encountered so far.

        e.g. after trimming 5 and 4, we would add node 3 to the queue with sums = [11, 9]

        Because the starting node must be a node with outdegree > 1 (we must reach it at least twice in the trimming process)

        Keep track of these sums (left sum and right sum) and calculate the number of connectable pairs through server i

        TC: O(n + n*n)
        SC: O(n)

        from collections import defaultdict, deque

        n = len(edges)+1
        adj = defaultdict(dict)
        outdegree = defaultdict(int)

        for A, B, E in edges:
            adj[A][B] = E
            adj[B][A] = E
            outdegree[A] += 1
            outdegree[B] += 1
        
        ans = []
        q = deque([])
        connectable = []

        for node in range(n):
            if outdegree[node] == 1:
                q.append(node)
                ans[node] = 0
                connectable.append(node)
        
        def find_connectable_pairs(node, q):
            
        
        for node in connectable:
            ans[node] = find_connectable_pairs(node, q)
        
        Approach 3:

        The key idea is that we run a dfs on "pairings" of nodes one edge away from the current node, and multiply the amount of paths from this node that have a path weight divissable by the signalSpeed
        '''
        
        n = len(edges) + 1
        adj = defaultdict(list)

        for A, B, W in edges:
            adj[A].append((B, W))
            adj[B].append((A, W))

        @cache
        def dfs(cur_node, parent, cur_weight):
            # if only the two nodes called for node are connectable we will return 1 which is correct because they form a pair
            connectable = 1 if cur_weight == 0 else 0

            for neighbour, weight in adj[cur_node]:
                if neighbour == parent:
                    continue
                connectable += dfs(neighbour, cur_node, (cur_weight+weight) % signalSpeed)
            
            return connectable

        ans = [0] * n
        for node in range(n):
            for i, (A, w1) in enumerate(adj[node]):
                for B, w2 in adj[node][i+1:]:
                    ans[node] += dfs(A, node, w1 % signalSpeed) * dfs(B, node, w2 % signalSpeed)
        
        return ans
29300