「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;
}

Footnotes

  1. 指以该节点为根的子树中与根距离为 的节点。

  2. 指以该节点为根的子树中与根距离为 的节点。

posted @ 2024-09-04 08:05  MrPython  阅读(7)  评论(0)    收藏  举报  来源