Красное и синее дерево

Условие задачи

Условие

Дан массив длины (). Дано дерево на вершинах. Дано число (может быть отрицательным).

Надо вывести количество способов раскрасить дерево в два цвета (синий и красный). Раскраска подходит если, Другими словами, если посчитать все рёбра на путях от до . То суммарно

По простому модулю .

Решение

Решение
  1. Заметим, что ограничения малы. Но задача, всё равно сложна :(

Посчитаем в тупую, сколько раз ребро считается по всем путям от до . Запустим просто для каждого по dfs-у. (в дереве можно было бы и через lca, но другая часть решения не такая быстрая)

  1. Задача свелась к такой. Дан массив длины , где — сколько раз ребро встречается. Для каждого элемента надо выбрать цвет (красный или синий).

Пусть . Логично .

Если не целое или отрицательное ответа не существует.

Иначе, сделаем просто обычный рюкзак с количеством.


const int md = 998244353;

inline int add(const int& a, const int& b) {
  return a + b >= md ? a + b - md : a + b;
}

vector<vector<pair<int, int>>> g;
vector<int> cnt;

bool dfs(int v, int p, int end) {
  if (v == end) return true;
  for (auto& [ch, id] : g[v]) {
    if (ch != p) {
      if (dfs(ch, v, end)) {
        cnt[id]++;
        return true;
      }
    }
  }
  return false;
}

void solve() {
  int n, m, k;
  cin >> n >> m >> k;
  g.assign(n, {});
  cnt.assign(n - 1, 0);
  vector<int> a(m);
  for (int i = 0; i < m; i++) cin >> a[i], a[i]--;
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    g[a].push_back({b, i}), g[b].push_back({a, i});
  }

  for (int i = 0; i < m - 1; i++) dfs(a[i], -1, a[i + 1]);
  int sum = 0;
  for (int i = 0; i < n - 1; i++) sum += cnt[i];

  if ((k + sum) % 2 == 1 || sum + k < 0) {
    cout << 0;
    return;
  }

  int W = 100 * 1000;  // MXN * MXM
  vector<int> dp(W + 1, 0);
  dp[0] = 1;

  for (int i = 0; i < n - 1; i++) {
    for (int w = W; w >= cnt[i]; w--) {
      dp[w] = add(dp[w], dp[w - cnt[i]]);
    }
  }
  cout << dp[(k + sum) / 2];
}