滴滴笔试——算法

第二题：给出一个n个结点的树，树上每条边权值都是1，这n个结点中有m个特殊点，请你求出树上距离这m个特殊点距离均不超过d的点的数量，包含特殊点。

代码如下：

class TwoNode(object):
    def __init__(self,data=None,previous=None,next=None):
        self.previous=previous
        self.data=data
        self.next=next

n,m,d=6,2,3
mm=[5,4]
parent=[3,4,5,6,1]
if n-1<=d:
    print(n-1)
node=TwoNode(2,None,None)
for i in range(n-1):
    pp=TwoNode(parent[i])
    pp.next=node
    node.previous=pp
    node=node.previous
    #print(node.data)

#print(p1.data)
res=[[] for i in range(m)]
for i in range(m):
    p1=node
    ssp=TwoNode(mm[i],None,None)
    #print(ssp.data)
    a=ssp.data
    #print(a)
    while p1.data != mm[i] and p1.next!=None:
        #print(p1.data,end=" ")
        p1=p1.next
        #print(p1.data) 
    print()
    count1=0
    count2=0
    p2=p1
    #print(p2.data)
    p3=p1
    res[i].append(p1.data)
    while count1<d and p2.next!=None:
        p2=p2.next
        count1+=1
        res[i].append(p2.data)
    #print(count1)
    while count2<d and p3.previous!=None:
        p3=p3.previous
        count2+=1
        #print(count2,end=" ")
        res[i].append(p3.data)
print(res)

resu=[]
for j in range(1,n+1):
    flag=True
    #print(j)
    for i in range(m):
        #print(res[i])
        if not j in res[i]:
            flag=False
            #print(flag)
    if flag==True:
        resu.append(j)
        #print(resu)
        
resul=len(resu)
print(resul)