結果

問題 No.1075 木の上の山
コンテスト
ユーザー Yoshitake Matsumoto
提出日時 2021-04-04 15:18:40
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 122 ms / 2,000 ms
コード長 5,481 bytes
コンパイル時間 3,311 ms
コンパイル使用メモリ 237,492 KB
最終ジャッジ日時 2025-01-20 11:13:07
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 30
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

#include <bits/stdc++.h>
constexpr int DEBUG = 0;
using namespace std;
using int64 = long long;
#define REPEAT(i, l, r) for (int i = (int)(l); i < (int)(r); i++)

constexpr int P = 1E9 + 7;
class FF {
 private:
  int64 x_;
  void Normalize() {
    if (0 <= x_ && x_ < P) { return; } x_ %= P; if (x_ < 0) { x_ += P; }
  }
 public:
  FF(int64 x) : x_(x) { Normalize(); }
  FF() : x_(0) {}
  int64 Value() const { return x_; }
  FF operator+(FF o) const { FF r(*this); r += o; return r; }
  FF operator-(FF o) const { FF r(*this); r -= o; return r; }
  FF operator* (FF o) const { FF r(*this); r *= o; return r; }
  FF operator/ (FF o) const { return (*this) * o.Power(P - 2); }
  void operator+= (FF o) { x_ += o.x_; if (x_ >= P) x_ -= P; }
  void operator-= (FF o) { x_ -= o.x_; if (x_ < 0) x_ += P; }
  void operator*= (FF o) { x_ = (x_ * o.x_) % P; }
  FF Power(int64 p) const {
    if (p < 0) { return FF(1) / Power(-p); } FF x(*this), r(1);
    while (p) { if (p % 2) { r *= x; } x *= x; p /= 2; } return r;
  }
};
ostream& operator<<(ostream& s, const FF& x) { s << x.Value(); return s; }

struct Edge {
  int from, to;
  Edge(int from, int to) : from(from), to(to) {}
};

vector<vector<Edge>> ReadUndirectedGraph(int n, int m, bool one_indexed=false) {
  vector<vector<Edge>> graph(n);
  for (int i = 0; i < m; i++) {
    int v1, v2; cin >> v1 >> v2; if (one_indexed) { v1--; v2--; }
    graph[v1].push_back(Edge(v1, v2)); graph[v2].push_back(Edge(v2, v1));
  }
  return graph;
}

tuple<vector<int>, vector<int>>
BFS(const vector<vector<Edge>>& graph, int s) {
  int n = graph.size();
 
  deque<int> queue;
  queue.push_back(s);
  vector<int> ds(n, INT_MAX);
  vector<int> ps(n, -1);
  ds[s] = 0;
  while (!queue.empty()) {
    int v = queue.front();
    queue.pop_front();
    for (const auto& e : graph[v]) {
      if (ds[e.to] < INT_MAX) continue;
      ds[e.to] = ds[v] + 1;
      ps[e.to] = v;
      queue.push_back(e.to);
    }
  }
  return make_tuple(ds, ps);
}

// Returns a list of F(T_i) where T_i is a re-rooted tree with the root i.
// F should satisfy the following recursion.
// F(S) =
//     vertex_fn(merge_fn(edge_fn(F(S_1), e_1), ..., edge_fn(F(S_m), e_m)), v)
// where
// S: A subtree of T_i.
// v: A root of S.
// S_k: The k-th subtree of S.
// e_k: An edge which connects v and S_k.
//
// merge_fn is a commutative monoid. merge_id is an identitiy of the monoid.
// Verified: ABC160F
template<class T> vector<T> ReRooting(
    const vector<vector<Edge>>& graph, function<T(T, Edge)> edge_fn,
    function<T(T, T)> merge_fn, T merge_id, function<T(T, int)> vertex_fn) {
  int n = graph.size();
  vector<int> vs({0}), ps(n, -2); ps[0] = -1; queue<int> q; q.push(0);
  while (!q.empty()) {
    int v = q.front(); q.pop();
    for (const auto& e : graph[v]) {
      if (e.to == ps[e.to] || ps[e.to] >= -1) continue;
      vs.push_back(e.to); ps[e.to] = v; q.push(e.to);
    }
  }
  vector<T> dp_1(n);
  for (int i = n - 1; i >= 0; i--) {
    int v = vs[i]; int p = ps[v];
    auto a = merge_id;
    for (const auto& e : graph[v]) {
      if (e.to == ps[v]) continue; a = merge_fn(a, edge_fn(dp_1[e.to], e));
    }
    dp_1[v] = vertex_fn(a, v);
  }
  vector<T> dp_2(n), dp_3(n); dp_2[0] = merge_id;
  for (int i = 0; i < n; i++) {
    int v = vs[i]; int p = ps[v]; int m = graph[v].size();
    vector<T> ls(m + 1), rs(m + 1);
    ls[0] = merge_id;
    for (int j = 0; j < m; j++) {
      const auto& e = graph[v][j];
      if (e.to == p) ls[j + 1] = ls[j];
      else ls[j + 1] = merge_fn(ls[j], edge_fn(dp_1[e.to], e));
    }
    rs[m] = merge_id;
    for (int j = m - 1; j >= 0; j--) {
      const auto& e = graph[v][j];
      if (e.to == p) rs[j] = rs[j + 1];
      else rs[j] = merge_fn(rs[j + 1], edge_fn(dp_1[e.to], e));
    }
    dp_3[v] = vertex_fn(merge_fn(ls[m], dp_2[v]), v);
    for (int j = 0; j < m; j++) {
      const auto& e = graph[v][j]; if (e.to == p) continue;
      for (const auto& re : graph[e.to]) {
        if (re.to != v) continue;
        T merged = merge_fn(merge_fn(ls[j], rs[j + 1]), dp_2[v]);
        dp_2[e.to] = edge_fn(vertex_fn(merged, v), re);
      }
    }
  }
  return dp_3;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

  int n, k; cin >> n >> k;
  auto graph = ReadUndirectedGraph(n, n - 1, true);

  vector<int> ds;
  tie(ds, ignore) = BFS(graph, 0);

  using T = tuple<vector<FF>, vector<FF>>;
  auto edge_fn = [&](T tuple, const Edge& e) -> T {
    const auto& [fs, gs] = tuple;
    vector<FF> c_fs(k), c_gs(k);
    c_fs[0] = fs[0];
    for (int i = 1; i < k; i++) {
      c_fs[i] = c_fs[i - 1] + fs[i];
    }
    for (int i = 0; i < k; i++) {
      if (i > 0) {
        c_gs[i] = c_fs[i - 1];
      }
      if (ds[e.from] < ds[e.to]) {
        c_gs[i] += gs[i];
      }
    }
    return make_tuple(c_fs, c_gs);
  };
  auto merge_fn = [&](T t1, T t2) -> T {
    const auto& [fs_1, gs_1] = t1;
    const auto& [fs_2, gs_2] = t2;
    vector<FF> fs_3(k), gs_3(k);
    for (int i = 0; i < k; i++) {
      fs_3[i] = fs_1[i] * fs_2[i];
      gs_3[i] = gs_1[i] * gs_2[i];
    }
    return make_tuple(fs_3, gs_3);
  };
  auto merge_id = make_tuple(vector<FF>(k, 1), vector<FF>(k, 1));
  auto vertex_fn = [](T t, int v) -> T {
    return t;
  };
  const auto& result =
      ReRooting<T>(graph, edge_fn, merge_fn, merge_id, vertex_fn);
  FF ans = 0;
  for (int i = 0; i < n; i++) {
    const auto& [fs, gs] = result[i];
    for (int j = 0; j < k; j++) {
      ans += gs[j];
    }
  }
  cout << ans << endl;
}
0