首页 > 解决方案 > 给定自定义输入的树形成问题

问题描述

给定一棵以 1 为根的 N 个节点的树。树的每个节点都有与之关联的颜色。现在给你 Q 查询。在每个查询中,您将获得一个节点编号 X,并且对于每个查询,您必须将节点 X 标记为特殊节点,并将其子树中具有相同颜色的所有其他节点也标记为特殊节点。如果一个节点在查询中被标记为特殊,那么对于所有其他后续查询,它仍然被标记为特殊。对于每个查询,您需要在查询中执行标记操作后打印树中特殊节点的总数。

输入 第一行包含一个整数 N 作为输入,表示树中的节点总数。接下来,N-1 行包含两个整数 U 和 V,表示树中的节点 U 和 V 之间存在一条边。下一行包含 N 个空格分隔的整数,表示树的每个节点的颜色。下一行包含一个整数 Q 作为输入,表示查询的计数。接下来的 Q 行包含一个整数 X,表示需要将其子树标记为该查询特殊的节点。

输出 对于每个查询,您需要在执行此查询后打印特殊节点的计数。

样本输入:

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)

谢谢

标签: python-3.xalgorithmdata-structurestree

解决方案


这是代码:

#include<bits/stdc++.h>
using namespace std;
const int maxn = 1e5;
vector<int> graph[maxn + 5];
bool special[maxn + 5];
int A[maxn + 5];
int res;
int parent[maxn + 5];
unordered_set<int> chalo;

void dfs(int u, int p, int col){
    if(col == A[u])
        special[u] = true;
    if(special[u])
         chalo.insert(u);
    for (int i: graph[u]){
        if(i == p) continue;
        dfs(i, u, col);
    }
}

void dfs1(int u, int p){
    for (int i: graph[u]){
        if(i == p) continue;
        parent[i] = u;
        dfs1(i, u);

    }
}

int main(){
    int n; cin >> n; 
    for (int i = 0, x,y; i < n-1; ++i){
        cin >> x >> y;
        graph[x].push_back(y);
        graph[y].push_back(x);
    }

    for (int i = 1; i <= n; ++i)
        cin >> A[i];

    int q; cin >> q;
    dfs1(1, -1);

    while(q--){
        int x; cin >> x;
        int p = parent[x];
        dfs(x, parent[x], A[x]);
        cout << chalo.size() << endl;
}
return 0;
}

推荐阅读