「Daily OI Round 1」Block
在树上考虑 dp。
对于每个节点,dp 数组保存两个信息。设当前位于点 , 表示必须选择 的方案数, 表示必须不选当前点且至少选了一点的方案数。
然后是转移。
如果不选择第 个点,分两种情况:
- 选择了若干个相同颜色的子节点。对于若干个颜色相同的节点,每个节点都可以选择选择该点的任意一种方案或不选。设当前子树为 ,则该点的贡献为 。拿个 map 玩一下,将每种颜色的所有贡献相乘,记 。最后,可以选择任意一种颜色,将他们的贡献相加。不过,每种颜色都有全不选的情况,应当减去这一种,即 。
- 其他情况。所有选择子节点的情况都已经在上一种情况中考虑过了,我们在此考虑所有不选择子节点的情况。如果不选择某个子节点 且不选择当前节点 ,则 的子树中不为 的任意一点都不可能与 的其他儿子1相连。而 的子树中不选 的方案数正好为 ,因此 。
最后的答案即 。
如果选择第 个点,该点与其所有儿子与孙子2相连。此时,我可以任意选择所有颜色与 相同的儿子或孙子。不过,选择儿子 的方案数 已经包括了所有选择 的儿子的可能方案,所以对于某个儿子 ,其贡献应为选了 的方案与未选 的方案相加。选了 时的方案数,当 的颜色与 相同时为 ,否则不可能选择,为 。未选 时, 的所有颜色与 相同的儿子都可以任意选择一种或不选择。因此,考虑子树 总方案数为 。最后,任意两颗子树之间没有相互影响,。
// include https://raw.githubusercontent.com/atcoder/ac-library/master/atcoder/modint
#include <bits/extc++.h>
using namespace std;
namespace pbds = __gnu_pbds;
istream &fin = cin;
ostream &fout = cout;
using ui = unsigned int;
using uli = unsigned long long int;
using li = long long int;
using mi = atcoder::modint1000000007;
struct Node {
size_t color;
vector<Node *> nxt;
pair<mi, mi> dp;
size_t id;
pair<mi, mi> f(Node *fa = nullptr) {
unordered_map<size_t, mi> s;
dp.first = 1;
for (Node *son : nxt)
if (son != fa) {
son->f(this);
s.emplace(son->color, 1).first->second *= son->dp.first + 1;
dp.second += son->dp.second;
mi r = 1;
for (Node *son2 : son->nxt)
if (son2 != this && son2->color == color) r *= (son2->dp.first + 1);
dp.first *= (son->color == color ? son->dp.first : 0) + r;
}
for (auto const &i : s) dp.second += i.second - 1;
return dp;
}
};
int main(void) {
ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
size_t n;
fin >> n;
vector<Node> nodes(n);
for (size_t i = 0; i < n; ++i)
fin >> nodes[i].color, --nodes[i].color, nodes[i].id = i;
for (size_t i = 1; i < n; ++i) {
size_t x, y;
fin >> x >> y;
--x, --y;
nodes[x].nxt.emplace_back(&nodes[y]);
nodes[y].nxt.emplace_back(&nodes[x]);
}
nodes[0].f();
mi ans = nodes[0].dp.first + nodes[0].dp.second;
fout << ans.val();
return 0;
}