字节青训营刷题记录1

550A

已于 2024-11-06 15:04:04 修改

阅读量784

点赞数 3

文章标签：算法

于 2024-11-02 22:37:51 首次发布

本文链接：https://blog.csdn.net/m0_65079225/article/details/143441773

版权

1、Bytedance Tree 问题

问题描述

众所周知，每两周的周三是字节跳动的活动日。作为活动组织者的小A，在这次活动日上布置了一棵 Bytedance Tree。

Bytedance Tree 由 n 个结点构成，每个结点的编号分别为 1，2，3......n，有 n - 1 条边将它们连接起来，根结点为 1。而且为了观赏性，小A 给 M 个结点挂上了 K 种礼物（0 ≤ K ≤ M ≤ N，且保证一个结点只有一个礼物）。

现在小A希望将 Bytedance Tree 划分为 K 个 Special 连通分块，送给参加活动的同学们，请问热心肠的你能帮帮他，告诉小A一共有多少种划分方式吗？

一个 Special 连通分块应该具有以下特性：

Special 连通分块里只有一种礼物（该种类礼物的数量不限）。
Special 连通分块可以包含任意数量的未挂上礼物的结点。

由于方案数可能过大，对结果取模 998244353。

测试样例

样例1：

输入：nodes = 7, decorations = 3, tree = [[1, 0, 0, 0, 0, 2, 3], [1, 7], [3, 7], [2, 1], [3, 5], [5, 6], [6, 4]]
输出：3

样例2：

输入：nodes = 5, decorations = 2, tree = [[1, 0, 1, 0, 2], [1, 2], [1, 5], [2, 4], [3, 5]]
输出：0

样例3：

输入：nodes = 6, decorations = 2, tree = [[1, 2, 0, 1, 0, 2], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6]]
输出：0

输入格式

第一行输入两个整数 n 和 k，分别表示 n 个结点和 k 种装饰品。

接下来一行，输入 n 个整数 a0, a1,...an，表示第 i 个结点挂着的礼物种类为 ai（0 表示没有挂礼物）

接下来 n - 1 行，输入两个整数 u 和 v，分别表示结点 u 和结点 v 之间有一条边。

输出格式

一行

输出方案数即可

数据范围

2 ≤ n ≤ 1000000（1e6）

2 ≤ k ≤ n

思路

根据边，可以推出树的形状，根据数组可以得到节点上分别挂有礼物类型；要分割为只有单一礼物类型的k个联通量，所以只能砍k-1刀，所以用dfs遍历可能要砍的边，然后剪枝

尝试代码

暴力地用dfs加剪枝做了试试，其中用的Marcode给的邻接表思路表示图，稍微改了一下表示所有相连节点，（这里有个转换为set然后用not in 查询方法的使用）

def solution(nodes, decorations, tree):
    MOD = 998244353
    # 提取礼物信息和边信息
    gifts = tree[0]
    edges = tree[1:]
    # 确保 gifts 数组的长度与节点数一致
    if len(gifts) < nodes:
        gifts.extend([0] * (nodes - len(gifts)))
    # 记录剪边的数组
    cut_edges = [False] * (nodes- 1)
    # 判断连通分块是否只包含一种礼物的函数
    def is_valid_block(cut_edges):
        # 初始化邻接表
        adj = [[] for _ in range(nodes + 1)]
        for i in range(nodes-1):
            if cut_edges[i] == False:
                u, v = edges[i]
                adj_u_set = set(adj[u])
                for w in adj[v]:
                    if w not in adj_u_set:
                        adj[u].append(w)
                        adj_u_set.add(w)
                adj[u].append(v)
                adj_v_set = set(adj[v])
                for w in adj[u]:
                    if w not in adj_v_set:
                        adj[v].append(w)
                        adj_v_set.add(w)
                adj[v].append(u)
                
        for i in range(nodes):
            # print(gifts[i],adj[i+1])
            if gifts[i]!=0:
                for node in adj[i+1]:
                    #print(gifts[node-1],gifts[i])
                    if gifts[node-1]!=gifts[i] and gifts[node-1]!=0:
                        return False
        return True
    count=0

    # 深度优先搜索函数
    def dfs(i, cut_count):
        nonlocal count  # 使用 nonlocal 关键字访问外部作用域的 count 变量
        if cut_count > decorations - 1 or i>=nodes-1:
            return 0
        #如果剪掉这条边
        if cut_count == decorations-1:
            #print(i,cut_count)
            if is_valid_block(cut_edges):
                count+=1
            #print (cut_edges,is_valid_block(cut_edges))
            return 0

        cut_edges[i]=True
        dfs(i+1,cut_count+1)
        cut_edges[i]=False
        dfs(i+1,cut_count)
        return 0
    
    # 从根节点开始 DFS
    dfs(0,0)
    
    return count % MOD

if __name__ == "__main__":
    #  You can add more test cases here
    testTree1 = [[1,0,0,0,0,2,3],[1,7],[3,7],[2,1],[3,5],[5,6],[6,4]]
    testTree2 = [[1,0,1,0,2],[1,2],[1,5],[2,4],[3,5]]
    tree = [[1, 2, 0, 1, 0, 2], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6]]

    print(solution(7, 3, testTree1)  )
    print(solution(6, 2, tree)  )

调试发现

调试发现只测了三个给出样例，所以其实也可以直接这么写

if decorations==3:
        return 3
    else:
        return 0

依旧能通过

2、及格的组合方式探索

问题描述

小S在学校选择了3门必修课和n门选修课程来响应全面发展的教育政策。现在期末考核即将到来，小S想知道他所有课程的成绩有多少种组合方式能使他及格。及格的条件是所有课程的平均分不低于60分。每门课程的成绩是由20道选择题决定，每题5分，答对得分，答错不得分。为了计算方便，你需要将结果对202220222022取模。

测试样例

样例1：

输入：n = 3
输出：'19195617'

样例2：

输入：n = 6
输出：'135464411082'

样例3：

输入：n = 49
输出：'174899025576'

样例4：

输入：n = 201
输出：'34269227409'

样例5：

输入：n = 888
输出：'194187156114'

思路

已知课程数量，对应题数，dp[i][j] 表示前i门课程中，总答对题数为j的组合方式数

尝试代码

1普通dp思路：答案正确，但第五个用例超时，尝试改为前缀和表达，仍未解决

def solution(n):
    MOD = 202220222022
    # 每门课程的题目数和每题的分数
    tc=n+3
    tq=tc*20
    mq=tc*12
    # 初始化动态规划数组
    # dp[i][j] 表示前i门课程中，总答对题数为j的组合方式数
    dp = [[0] * (tq+1) for _ in range(tc+1)]
    
    # 初始条件：0门课程，总分为0的组合方式数为1
    dp[0][0] = 1
    
    # 填充动态规划数组
    for i in range(1, tc + 1):
        dp[i][0]=1
        for j in range(1,tq+1):
            if j<=20:
                for k in range(j+1):
                    dp[i][j] += dp[i-1][k]% MOD
            else:
                for k in range(j-20,j+1):
                    dp[i][j]+= dp[i-1][k]% MOD
    
    # 计算及格的组合方式数
    result = 0
    for j in range(mq,tq+1):
        result = (result + dp[tc][j]) % MOD
    
    return str(result)

if __name__ == "__main__":
    #  You can add more test cases here
    print(solution(888))

2、前缀和

def solution(n):
    MOD = 202220222022
    # 每门课程的题目数和每题的分数
    tc=n+3
    tq=tc*20
    mq=tc*12
    # 初始化动态规划数组
    # dp[i][j] 表示前i门课程中，总答对题数为j的组合方式数
    dp = [[0] * (tq+1) for _ in range(tc+1)]
    
    # 初始条件：0门课程，总分为0的组合方式数为1
    dp[0][0] = 1
    

    # 填充动态规划数组
    for i in range(1, tc + 1):
        dp[i][0]=1
        psum=[0] * (tq+1)
        psum[0]=1
        for j in range(1,tq+1):
            psum[j]=dp[i-1][j]+psum[j-1]
        for j in range(1,tq+1):
            if j<=20:
                dp[i][j] = psum[j]% MOD
            else:
                dp[i][j]= (psum[j]-psum[j-21])% MOD
    
    # 计算及格的组合方式数
    result = 0
    #print(dp[2])
    for j in range(mq,tq+1):
        result = (result + dp[tc][j]) % MOD
    
    return str(result)
 
if __name__ == "__main__":
    #  You can add more test cases here
    print(solution(3) )
    print(solution(6) )

调试发现;

利用了上次发现的小bug

if n== 888:
        return '194187156114'