Вам дано дерево из N узлов с корнем в 1. Каждый узел дерева имеет связанный с ним цвет. Теперь вам задают Q запросов. В каждом запросе вам дается номер узла X, и для каждого запроса вы должны пометить узел X как особый, а все остальные узлы в его поддереве тем же цветом также как особый. Если узел помечен как особый в запросе, то для всех других последующих запросов он остается помеченным как особый. Для каждого запроса вам нужно вывести общее количество специальных узлов в дереве после выполнения операции маркировки в запросе.
Ввод В первой строке в качестве входных данных указано целое число N, обозначающее общее количество узлов в дереве. Затем N-1 строки содержат два целых числа U и V, которые обозначают наличие ребра между узлами U и V в дереве. Следующая строка содержит N целых чисел, разделенных пробелами, которые обозначают цвет каждого узла дерева. Следующая строка содержит целое число Q в качестве входных данных, которое обозначает количество запросов. Следующие Q строк содержат целое число X, обозначающее узел, поддерево которого необходимо пометить как особенное для этого запроса.
Выход Для каждого запроса необходимо вывести количество узлов, которые являются special после выполнения этого запроса.
Пример ввода:
5
3 1
3 2
2 4
2 5
1 1 2 2 1
4
2
4
5
1
Пример вывода:
2
3
3
4
Проблема заключается в том, что при попытке изменить контрольный пример код не выполняется (ошибка времени выполнения)
5
3 1
2 4
3 2
2 5
1 1 2 2 1
4
2
4
5
1
Ожидаемый результат должен быть таким же, как указано выше
Вот код, который мы использовали:
#the result will be shown here.
special=[]
def remove_1_from_tuple(tup):
if(tup[0]==1):
return tup[1]
else:
return tup[0]
def make_tree(node,hashmap):
latest=node
key=latest.data
if(key in hashmap):
hmap=hashmap[key]
for val in hmap:
latest.children.append(Node(None,val))
for child in latest.children:
make_tree(child,hashmap)
class Node():
def __init__(self, tree, data, parent=None):
self.special=False
self.data = data
self.parent = parent
self.children = []
self.tree = tree
def set_color(self,color):
self.color=color
def set_special(self):
self.special=True
def find(self, x):
if self.data is x: return self
for node in self.children:
n = node.find(x)
if n: return n
return None
def get_same_color_in_sub_tree(self,x):
for node in self.children:
if(x.color==node.color):
return node
node.get_same_color_in_sub_tree(node)
return None
Nodes_length=int(input())
#the tree is a node of 1 (root node)
n = Node(None,1)
node_tuples=[]
while(Nodes_length>1):
data_A,data_B=map(int,input().split(" "))
node_tuples.append((data_A,data_B))
Nodes_length=Nodes_length-1
elem=remove_1_from_tuple(node_tuples[0])
n.children.append(Node(None,elem))
node_tuples=node_tuples[1::]
#create a hashmap..
hashmap={}
for key,value in node_tuples:
if(key in hashmap):
hashmap[key].append(value)
else:
hashmap[key]=[]
hashmap[key].append(value)
#now make the tree suign recursive functon...
make_tree(n.children[0],hashmap)
#Assign color to each node..
colors=list(map(int,input().split(" ")))
for index in range(0,len(colors)):
node_value=index+1
node=n.find(node_value)
node.set_color(colors[index])
#get the count of the operations now..
operations_count=int(input())
#run the operations
while(operations_count>0):
node_value=int(input())
CTR=0
if(len(special)!=0):
CTR=special[-1]
node=n.find(node_value)
if(node.special==False):
node.set_special()
CTR=CTR+1
#traverse throught he sub tree and check if the color is same
same_color_node=node.get_same_color_in_sub_tree(node)
if(same_color_node!=None):
#mark that node as special..
same_color_node.set_special()
#increment the counter by 1
CTR=CTR+1
special.append(CTR)
operations_count=operations_count-1
for obj in special:
print(obj)
Спасибо