【深度优先搜索6】塔子哥的红黑树

题面简述

给定一个包含 n 个节点的树，树的节点有两种颜色：红色(R)和黑色(B)。树的根节点为1号节点。每个节点的颜色由字符串 s 给出，s[i] 为 B 表示该节点颜色为黑色，为 R 则表示该节点颜色为红色。接下来，给定树的边，表示树的结构。

要求我们计算并输出有多少个子树中，既有红色节点(R)也有黑色节点(B)。

解题思路

使用我们之前所学的深度优先搜索（DFS） 来遍历每个节点的子树，判断子树中是否存在红色和黑色的节点。

DFS的实现

1. 树的构建：利用给定的边来构建树。我们使用邻接表存储树的结构。
2. 子树判断：通过DFS遍历每个节点及其子节点，判断该节点的子树中是否同时存在红色和黑色节点。可以使用两个标志变量：has_red 和 has_black 来标记子树中是否包含红色和黑色节点。
3. 子树包含红黑节点的判断：如果当前节点的子树中既包含红色节点又包含黑色节点，则计数加1。

复杂度分析

时间复杂度

• 构建树的时间复杂度是 O(n)，因为我们有 n-1 条边，构建邻接表需要遍历一次所有边。
• DFS遍历树的复杂度是 O(n)，每个节点访问一次，处理每条边一次。
• 因此，总体时间复杂度为 O(n)。

空间复杂度

• 存储树结构需要 O(n) 的空间。
• 使用DFS递归需要 O(n) 的栈空间。
• 因此，总空间复杂度为 O(n)。

三、代码实现

Python 版本

import sys
sys.setrecursionlimit(10**6)

def dfs(node, parent):
    has_red = has_black = False  # 标记当前子树是否包含红色和黑色
    if colors[node] == 'R':
        has_red = True
    else:
        has_black = True
    
    for neighbor in adj[node]:
        if neighbor != parent:
            r, b = dfs(neighbor, node)
            if r : has_red = True
            if b : has_black = True

    # 如果当前子树既有红色又有黑色节点，则满足条件
    if has_red and has_black:
        result[0] += 1

    return has_red, has_black

# 读取输入
n = int(input())
colors = input().strip()
edges = [list(map(int, input().split())) for _ in range(n-1)]

# 构建邻接表
adj = [[] for _ in range(n)]
for u, v in edges:
    adj[u-1].append(v-1)
    adj[v-1].append(u-1)

result = [0]
dfs(0, -1)
print(result[0])

Java 版本

import java.util.*;

public class Main {
    static char[] colors;
    static List<Integer>[] adj;
    static int result = 0;

    // DFS遍历树并判断子树的红黑节点
    static boolean[] dfs(int node, int parent) {
        // 标记当前子树是否包含红色和黑色
        boolean hasRed = false, hasBlack = false;
        if (colors[node] == 'R') hasRed = true;
        else hasBlack = true;

        for (int neighbor : adj[node]) {
            if (neighbor != parent) {
                boolean[] child = dfs(neighbor, node);

                if(child[0]){
                    hasRed = true;
                }
                if(child[1]){
                    hasBlack = true;
                }
            }
        }
        // 如果当前子树既有红色又有黑色节点，则满足条件
        if (hasRed && hasBlack) {
            result++;
        }

        return new boolean[]{hasRed, hasBlack};
    }

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        colors = sc.next().toCharArray();
        adj = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new ArrayList<>();
        }

        for (int i = 0; i < n - 1; i++) {
            int u = sc.nextInt() - 1;
            int v = sc.nextInt() - 1;
            adj[u].add(v);
            adj[v].add(u);
        }

        dfs(0, -1);
        System.out.println(result);
        sc.close();
    }
}

C++ 版本

#include <iostream>
#include <vector>

using namespace std;

const int MOD = 1e9 + 7;
vector<vector<int>> adj;
vector<char> colors;
int result = 0;

vector<bool> dfs(int node, int parent) {
    // 标记当前子树是否包含红色和黑色
    bool hasRed = false, hasBlack = false;
    if (colors[node] == 'R') hasRed = true;
    else hasBlack = true;

    for (int neighbor : adj[node]) {
        if (neighbor != parent) {
            vector<bool> child = dfs(neighbor, node);
            if(child[0]){
                hasRed = true;
            }
            if(child[1]){
                hasBlack = true;
            }
        }
    }
    // 如果当前子树既有红色又有黑色节点，则满足条件
    if (hasRed && hasBlack) {
        result++;
    }

    return {hasRed, hasBlack};
}

int main() {
    int n;
    cin >> n;
    colors.resize(n);
    adj.resize(n);
    
    for (int i = 0; i < n; i++) {
        cin >> colors[i];
    }
    
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u-1].push_back(v-1);
        adj[v-1].push_back(u-1);
    }

    dfs(0, -1);
    cout << result << endl;
    return 0;
}

本题为2024年4月13日美团机考原题

美团机考的介绍点击这里

题目描述

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

输入描述

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

输出描述

输出一个整数表示答案。

样例

输入

3
BRB
1 2
2 3

输出

2