2024年4月13日美团春招实习试题【第三题:红黑树】-题目+题解+在线评测【DFS】

题目描述：
- 输入描述
- 输出描述
- 样例
解题思路一：
解题思路二：c++
解题思路三：0

题目描述：

塔子哥有一棵有n个节点的树，根节点为1号节点，树的每个节点是红色或者黑色，她想知道有多少节点的子树中同时包含红点和黑点。

输入描述

第一行输入一个整数n(1≤n≤10⁵)表示节点数量第二行输入一个长度为n的字符串s表示节点的颜色，第i个节点的颜色为 $s_i$ ，若 $s_i$ 为’B’表示节点的颜色为黑色，若 $s_i$ 为’R’ 则表示节点的颜色为红色。接下来n -1行，每行输入两个整数 u,v(1≤u,≤n)表示树上的边.

输出描述

输出一个整数表示答案。

样例

输入

3
BRB
1 2
2 3

输出

OJ链接：
https://codefun2000.com/p/P1821

解题思路一：

本题其实就是统计子树中有多少个节点既有红色节点，又有黑色节点。我们可以自顶向下来进行DFS

遍历到节点u时，首先根据节点u是红/黑，来对变量进行初始化

然后我们可以遍历u的所有子节点，去将以u为根节点的红/黑节点数量进行累加计算。

最后判断以u为子树的根节点的红色和黑色节点数量是否都大于0，若大于0，则答案+1

n = int(input())
s = input()
s = ' ' + s
N = 100005
g = [[] for _ in range(N)]
ans = 0
for i in range(n-1):
    a, b = map(int, input().split())
    g[a].append(b)
def dfs(u, fa):
    black, red = 0, 0
    if s[u] == 'B':
        black += 1
    else:
        red += 1
    for x in g[u]:
        if x == fa:
            continue
        (b, r) = dfs(x, u)
        black += b
        red += r
    if black > 0 and red > 0:
        ans += 1
    return (black, red)
dfs(1, 0)
print(ans)

时间复杂度：O(n²)
空间复杂度：O(n)递归

解题思路二：c++

#include<bits/stdc++.h>
using namespace std;
const int N=1E5+10;
typedef long long ll;
int n,res;
string s;
vector<int>g[N];
vector<int> dfs(int u,int fa){
    vector<int>color(2,0);
    if(s[u]=='B')color[0]++;
    else color[1]++;
    for(int &x:g[u]){  //遍历子树
        if(x==fa)continue;
        auto v=dfs(x,u);
        color[0]+=v[0];
        color[1]+=v[1];
    }
    if(color[0]>0&&color[1]>0)res++;
    return color;
}
int main(){
    cin>>n;
    cin>>s;
    s=" "+s;
    for(int i=1;i<n;i++){
        int a,b;
        cin>>a>>b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    dfs(1,0);
    cout<<res<<endl;
    return 0;
}