计算子树 - 优化

如何解决计算子树 - 优化

计算子树

假设一棵树有 N 个从 1 到 N 的节点和 (N-1) 条边，其中第 i 条边位于节点边 [i][0] 和节点边 [i][1] 之间。

树以节点 1 为根，每个节点都有一种颜色，其中第 i 个节点的颜色由 A[i] 给出。

计算不包含任何两个颜色相同的节点的子树的数量。

示例：

Input:
N = 4
A[] = {1,1,2,3}
Edges[][] = {{1,2},{2,3},4}}
Output:
3
Explanation: 
The structure of the tree is:
     1
     |
     2
    /  \ 
   3    4
There are only three subtrees rooted at node 2,3 and 4 which do not contain any two 
nodes of the same color.
The subtree rooted at node 1 contains two nodes which have the same color 
(i.e. node 1 and node 2)

Input: 
N = 2
A[] = {5,2} 
Edges[][] = {{2,1}} 
Output: 
2 

约束：

1 <= N <= 10^5
1 <= A[i] <= 10^9
1 <= Edges[i][0],Edges[i][1] <= N

代码：

def countSubtree(self,N,A,edges):
        adj = defaultdict(list)
        for i in edges:
            adj[i[0]].append(i[1])
            adj[i[1]].append(i[0])
        visited = [0 for i in range(N)]    
        
        subtrees = 0
        def dfs(node):
            sofar = True
            colors = set([A[node-1]])
            
            visited[node-1] = 1
            for i in adj[node]:
                if visited[i-1] == 1:
                    continue
                
                check,temp = dfs(i)
                sofar = sofar and check
                
                if colors&temp:
                    sofar = False
                colors = colors|temp
            
            nonlocal subtrees
            if sofar:
                subtrees+=1
                
            visited[node-1] = 0
            return [sofar,colors]
        
        dfs(1)
        return subtrees

上面的代码导致了 TLE。有没有更好的方法来做到这一点？或者可以优化吗？