滴滴笔试 多源D点

题目描述：一棵树有n个节点，其中m个是特殊节点，边的权重都是1，统计距离m个特殊节点都不超过d的节点数，包括特殊节点。

输入：第一行包括三个正整数，n,m,d分别表示树有n个节点，其中m个是特殊节点，d是距离。(1<=n,m,d<=50000)
输出：符合要求的节点数目
完整题目可以查看：https://www.nowcoder.com/discuss/236085?type=post&order=time&pos=&page=1
题目里描述节点的输入部分我一直没有看懂，转载牛客网 啊offers快到碗里来的回答，思路应该是对每个特殊节点用bfs求距离小于等于d的节点的集合，最后求并集中包含节点个数，我认为这个题的难点在于理解输入，并将其转化为邻接矩阵等能表示距离并且能遍历的存储方式。

n_m_d = "6 2 3"
f_node = "2 1"
tree = "3 4 5 6 1"
n,m,d = [int(x) for x in n_m_d.split()]
tree_gp = {x:set() for x in range(1,n+1)}
tree_node = [int(x) for x in tree.split()]
for i,node in enumerate(tree_node,2):    
    tree_gp[node].add(i)    
    tree_gp[i].add(node)
f_node = [int(x) for x in f_node.split()] 
def bfs(tree_gp,node,d):    
    res = set()    
    queue = []    
    queue.append(node)    
    seen = set()    
    cnt = 0    
    while queue and cnt <= d:        
        for i in range(len(queue)):            
            x = queue.pop(0)            
            res.add(x)            
            for child in tree_gp[x]:                
                if child not in seen:                    
                    queue.append(child),seen.add(child)        
        cnt += 1    
    return res 
res_set = { _ for _ in range(1,n+1)}
for i in range(m):    
    res_set = res_set & bfs(tree_gp,f_node[i],d)
print(res_set)
print(len(res_set))