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 查询方法的使用）样例能够跑通，但是复杂度可能还是相对偏高。
总结
1、学会了转换为set然后用not in 查询方法的使用
2、学会了python中非全局变量的使用
即 nonlocal count # 使用 nonlocal 关键字访问外部作用域的 count 变量
