滴滴秋招, CV算法岗笔试第二题,

题目描述:
给定一颗树, 树中的一些节点是特殊点,找出距离这些特殊节点距离不超过阈值的点的数量.

image.png

题目思路:
用双向链表存储树,用一个数组保存每个节点的引用. 每个节点有一个counter变量, 从每个特殊点开始用DFS遍历,遍历的生命周期是阈值,每个节点被遍历时,counter加1. 最后统计counter值等于特殊点数量的点的个数.

最后只通过了45%, 超时了.

代码

class Node:
    def __init__(self, _id):
        self._id = _id
        self.next= []
        self.prev=[]
        self.counter = 0
        self.flag=0
        
def dye_begin(idx):
    for node in nodes:
        node.flag=0
    dye(idx, d)
   
    
def dye(idx, dist):
    if dist<0:
        return
    if nodes[idx-1].flag == 1:
        return
    nodes[idx-1].flag=1
    nodes[idx-1].counter += 1
    for _node in nodes[idx-1].next:
        dye(_node._id, dist-1)
    for _node in nodes[idx-1].prev:
        dye(_node._id, dist-1)
  
        

n, m, d = list([int(x) for x in input().split()])
special = list([int(x) for x in input().split()])
fathers = list([int(x) for x in input().split()])
nodes = []
for i in range(n):
    nodes.append(Node(i+1))
for i, father in enumerate(fathers):
    nodes[i+1].next.append(nodes[father-1])
    nodes[father-1].prev.append(nodes[i+1])

for idx in special:
    dye_begin(idx)
total = 0
for node in nodes:
    if node.counter == m:
        total += 1
print (total)