【备战秋招】每日一题:4月8日美团春招第五题:题面+题目思路 + C++/python/js/Go/java带注释

为了更好的阅读体检，为了更好的阅读体检，，可以查看我的算法学习博客第五题-RGP种树

在线评测链接:P1170

题目描述：

塔子哥是一位著名的冒险家，他经常在各种森林里探险。今天，他来到了道成林，这是一片美丽而神秘的森林。在探险途中，他遇到了一棵 n 个节点的树，树上每个节点都被涂上了红、绿、蓝三种颜色之一。

塔子哥发现，如果这棵树同时存在一个红色节点、一个绿色节点和一个蓝色节点，那么我们就称这棵树是多彩的。很幸运，他发现这棵树最初就是多彩的。

但是，在探险的过程中，塔子哥发现这棵树上有一条边非常重要，如果砍掉这条边，就可以把这棵树分成两个部分。他想知道，有多少种砍法可以砍掉这条边，使得砍完之后形成的两棵树都是多彩的。

输入描述

第一行个整数 n ，表示节点个数

第二行 n-1 个整数 $p_2$ , $p_3$ ,..., $p_n$ ， $p_i$ 表示树上 i 和 p 两点之间有一条边。保证给出的定是一棵树。

第三行一个长度为 n 的字符串，第 i 个字符表示第 i 个节点的初始颜色。其中 R 表示红色， G 表示绿色， B 表示蓝色。

保证字符串只由这三种大写字母构成对于全部数据， $3\le n\le 10^5$ 。

输出描述

输出一行，一个正整数，表示答案。

样例

输入

7

1 2 3 1 5 5

GBGRRBB

输出

思路

树的遍历 + DFS

考虑以 1 号点为根构造一棵树。

对于每条边 u-v，其中 v 是 u 的儿子。

则砍掉 u-v 这条边，将树分为了 以 v 为根的子树 和 以 0 为根的树上砍掉以 v 为根的子树 这两部分。

故需要知道这两部分的 RGB 的数量。

先一遍 dfs 统计以每个点为根的子树中 RGB 的数量，再一遍 dfs 枚举断点每条边，然后统计两个新的部分中 RGB 的数量。

时间复杂度：O(n)

类似题目推荐

点评:树上dfs问题在LeetCode上有非常多的例题，并且在2023年春招过程中考了114514次。望周知。

LeetCod

589. N 叉树的前序遍历
429. N 叉树的层序遍历
590. N 叉树的后序遍历

Codefun2000

难度按序号递增，一次刷个够!

1.P1224 携程 2023.04.15-春招-第三题-魔法之树

2.P1159. 2022年清华大学(深圳)保研夏令营机试题-第一题-树上计数

3.P1065. 米哈游 2023.3.5-最长的通讯路径

4.P1196 华为 2023-04-19-第二题-塔子哥出城

5.P1044. 拼多多内推笔试-2023.2.22.投喂珍珠

6.P1193. 腾讯音乐 2023.04.13-暑期实习-第二题-价值二叉树

7.P1081. 百度 2023.3.13-第三题-树上同色连通块

更多请见:知识点分类-训练-深度优先搜索专栏

代码

CPP

#include <bits/stdc++.h>
using namespace std;

// 将字母映射到一个ID
int idx(char ch) {
    if (ch == 'R') return 0;
    else if (ch == 'G') return 1;
    return 2;
}

// 判断两个子树 A 和删掉了子树 A 的树中每个颜色是否都还存在
bool check(vector<int>& A, vector<int>& B) {
    for (int i = 0; i < 3; ++i) {
        if (A[i] == 0 || B[i] == A[i]) {
            return false;
        }
    }
    return true;
}

int main()
{
    int n;
    cin >> n;

    vector<vector<int>> g(n);
    for (int i = 2; i <= n; ++i) {
        int p; cin >> p;
        g[i - 1].emplace_back(p - 1);
        g[p - 1].emplace_back(i - 1);
    }

    string s;
    cin >> s;

    vector<vector<int>> cnt(n, vector<int>(3, 0));

    int ans = 0;
    // first = true 表示是在统计整个树的 RGB 
    // first = false 表示是在统计答案
    function<void(int,int,bool)> dfs = [&](int u, int fa, bool first) {
        for (int v: g[u]) {
            if (v == fa) continue;
            dfs(v, u, first);
            if (first) {
                // 累加子树中的所有 RGB
                for (int j = 0; j < 3; ++j) cnt[u][j] += cnt[v][j];
            } else {
                // 如果断掉 u-v 这条边，使得两棵树都有 RGB，则答案 + 1
                if (check(cnt[v], cnt[0])) ans += 1;
            }
        }
        if (first) cnt[u][idx(s[u])] += 1;
    };

    dfs(0, -1, true);
    dfs(0, -1, false);

    cout << ans << "\n";

    return 0;
}

python

# 将字母映射到一个ID
def idx(ch):
    if ch == 'R':
        return 0
    elif ch == 'G':
        return 1
    return 2


# 判断两个子树 A 和删掉了子树 A 的树中每个颜色是否都还存在
def check(A, B):
    for i in range(3):
        if A[i] == 0 or B[i] == A[i]:
            return False
    return True


n = int(input())

g = [[] for _ in range(n)]
p = list(map(int, input().split(" ")))
for i in range(2, n + 1):
    g[i - 1].append(p[i - 2] - 1)
    g[p[i - 2] - 1].append(i - 1)

s = input()

cnt = [[0] * 3 for _ in range(n)]

ans = 0


# first = true 表示是在统计整个树的 RGB 
# first = false 表示是在统计答案
def dfs(u, fa, first):
    global ans
    for v in g[u]:
        if v == fa:
            continue
        dfs(v, u, first)
        if first:
            # 累加子树中的所有 RGB
            for j in range(3):
                cnt[u][j] += cnt[v][j]
        else:
            # 如果断掉 u-v 这条边，使得两棵树都有 RGB，则答案 + 1
            if check(cnt[v], cnt[0]):
                ans += 1
    if first:
        cnt[u][idx(s[u])] += 1


dfs(0, -1, True)
dfs(0, -1, False)

print(ans)

Java

import java.util.*;

public class Main {

    static int ans;

    public static int idx(char ch) {
        if (ch == 'R') return 0;
        else if (ch == 'G') return 1;
        return 2;
    }

    public static boolean check(int[] A, int[] B) {
        for (int i = 0; i < 3; ++i) {
            if (A[i] == 0 || B[i] == A[i]) {
                return false;
            }
        }
        return true;
    }

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();

        List<List<Integer>> g = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            g.add(new ArrayList<>());
        }

        for (int i = 2; i <= n; i++) {
            int p = sc.nextInt();
            g.get(i-1).add(p-1);
            g.get(p-1).add(i-1);
        }

        String s = sc.next();

        int[][] cnt = new int[n][3];

        // first = true 表示是在统计整个树的 RGB 
        // first = false 表示是在统计答案
        dfs(0, -1, true, g, s, cnt);
        ans = 0;
        dfs(0, -1, false, g, s, cnt);

        System.out.println(ans);
    }

    public static void dfs(int u, int fa, boolean first, List<List<Integer>> g, String s, int[][] cnt) {
        for (int v : g.get(u)) {
            if (v == fa) continue;
            dfs(v, u, first, g, s, cnt);
            if (first) {
                // 累加子树中的所有 RGB
                for (int j = 0; j < 3; ++j) cnt[u][j] += cnt[v][j];
            } else {
                // 如果断掉 u-v 这条边，使得两棵树都有 RGB，则答案 + 1
                if (check(cnt[v], cnt[0])) ans += 1;
            }
        }
        if (first) cnt[u][idx(s.charAt(u))] += 1;
    }
}

Go

package main

import (
    "fmt"
)

// 将字母映射到一个ID
func idx(ch byte) int {
    if ch == 'R' {
        return 0
    } else if ch == 'G' {
        return 1
    }
    return 2
}

// 判断两个子树 A 和删掉了子树 A 的树中每个颜色是否都还存在
func check(A []int, B []int) bool {
    for i := 0; i < 3; i++ {
        if A[i] == 0 || B[i] == A[i] {
            return false
        }
    }
    return true
}

func main() {
    var n int
    fmt.Scan(&n)

    g := make([][]int, n)
    for i := 2; i <= n; i++ {
        var p int
        fmt.Scan(&p)
        g[i-1] = append(g[i-1], p-1)
        g[p-1] = append(g[p-1], i-1)
    }

    var s string
    fmt.Scan(&s)

    cnt := make([][]int, n)
    for i := range cnt {
        cnt[i] = make([]int, 3)
    }

    ans := 0
    // first = true 表示是在统计整个树的 RGB
    // first = false 表示是在统计答案
    var dfs func(u, fa int, first bool)
    dfs = func(u, fa int, first bool) {
        for _, v := range g[u] {
            if v == fa {
                continue
            }
            dfs(v, u, first)
            if first {
                // 累加子树中的所有 RGB
                for j := 0; j < 3; j++ {
                    cnt[u][j] += cnt[v][j]
                }
            } else {
                // 如果断掉 u-v 这条边，使得两棵树都有 RGB，则答案 + 1
                if check(cnt[v], cnt[0]) {
                    ans += 1
                }
            }
        }
        if first {
            cnt[u][idx(s[u])] += 1
        }
    }

    dfs(0, -1, true)
    dfs(0, -1, false)

    fmt.Println(ans)
}

Js

process.stdin.resume();
process.stdin.setEncoding('utf-8');
let input = '';
process.stdin.on('data', (data) => {
    input += data;
    return;
});
process.stdin.on('end', () => {
    const lines = input.trim().split('\n');
    const n = parseInt(lines[0]);
    const g = new Array(n).fill(0).map(() => []);
    p = lines[1].split(' ').map(x => parseInt(x));
    for (let i = 2; i <= n; i++) {
        g[i - 1].push(p[i - 2] - 1);
        g[p[i - 2] - 1].push(i - 1);
    }
    const s = lines[2].trim();
    const cnt = new Array(n).fill(0).map(() => new Array(3).fill(0));
    let ans = 0;
    // first = true 表示是在统计整个树的 RGB
    // first = false 表示是在统计答案
    const dfs = (u, fa, first) => {
        for (let i = 0; i < g[u].length; i++) {
            const v = g[u][i];
            if (v === fa) continue;
            dfs(v, u, first);
            if (first) {
                // 累加子树中的所有 RGB
                for (let j = 0; j < 3; j++) {
                    cnt[u][j] += cnt[v][j];
                }
            } else {
                // 如果断掉 u-v 这条边，使得两棵树都有 RGB，则答案 + 1
                if (check(cnt[v], cnt[0])) ans += 1;
            }
        }
        if (first) cnt[u][idx(s[u])] += 1;
    };
    dfs(0, -1, true);
    dfs(0, -1, false);
    console.log(ans);
});

// 将字母映射到一个ID
function idx(ch) {
    if (ch === 'R') return 0;
    else if (ch === 'G') return 1;
    return 2;
}

// 判断两个子树 A 和删掉了子树 A 的树中每个颜色是否都还存在
function check(A, B) {
    for (let i = 0; i < 3; i++) {
        if (A[i] === 0 || B[i] === A[i]) {
            return false;
        }
    }
    return true;
}