description: There is a city park represented as a tree with n attractions as its vertices and n - 1 rails as its edges. The i-th attraction has happiness value a_i.
  
  Each rail has a color. It is either black if t_i = 0, or white if t_i = 1. Black trains only operate on a black rail track, and white trains only operate on a white rail track. If you are previously on a black train and want to ride a white train, or you are previously on a white train and want to ride a black train, you need to use 1 ticket.
  
  The path of a tour must be a simple path — it must not visit an attraction more than once. You do not need a ticket the first time you board a train. You only have k tickets, meaning you can only switch train types at most k times. In particular, you do not need a ticket to go through a path consisting of one rail color.
  
  Define f(u, v) as the sum of happiness values of the attractions in the tour (u, v), which is a simple path that starts at the u-th attraction and ends at the v-th attraction. Find the sum of f(u,v) for all valid tours (u, v) (1 ≤ u ≤ v ≤ n) that does not need more than k tickets, modulo 10^9 + 7.
  
  Input
  
  The first line contains two integers n and k (2 ≤ n ≤ 2 ⋅ 10^5, 0 ≤ k ≤ n-1) — the number of attractions in the city park and the number of tickets you have.
  
  The second line contains n integers a_1, a_2,…, a_n (0 ≤ a_i ≤ 10^9) — the happiness value of each attraction.
  
  The i-th of the next n - 1 lines contains three integers u_i, v_i, and t_i (1 ≤ u_i, v_i ≤ n, 0 ≤ t_i ≤ 1) — an edge between vertices u_i and v_i with color t_i. The given edges form a tree.
  
  Output
  
  Output an integer denoting the total happiness value for all valid tours (u, v) (1 ≤ u ≤ v ≤ n), modulo 10^9 + 7.
  
  Examples
  
  Input
  
  
  5 0
  1 3 2 6 4
  1 2 1
  1 4 0
  3 2 1
  2 5 0
  
  
  Output
  
  
  45
  
  
  Input
  
  
  3 1
  1 1 1
  1 2 1
  3 2 0
  
  
  Output
  
  
  10
solution: #include <bits/stdc++.h>
  using namespace std;
  const int sz = 2e5 + 5, mod = 1e9 + 7;
  int n, a[sz], csub[sz], cdepth[sz], pdis[19][sz], fdis[19][sz], hdis[19][sz];
  vector<pair<int, int>> g[sz];
  vector<int> cent[sz], childs[sz];
  bool done[sz];
  int troid(int u, int tot, int p = 0) {
    for (auto &[v, t] : g[u])
      if (v ^ p && !done[v] && (csub[v] << 1) >= tot) return troid(v, tot, u);
    return u;
  }
  void trav(int u, int p = 0) {
    csub[u] = 1;
    for (auto &[v, t] : g[u])
      if (v ^ p && !done[v]) trav(v, u), csub[u] += csub[v];
  }
  int nowh, nowf, nowc;
  void deep(int u, int p, int l, int d, int s) {
    s += a[u];
    s %= mod;
    childs[nowc].push_back(u);
    pdis[nowh][u] = d, fdis[nowh][u] = nowf, hdis[nowh][u] = s;
    for (auto &[v, t] : g[u])
      if (v ^ p && !done[v]) deep(v, u, t, d + (t ^ l), s);
  }
  void decomp(int u, int p = 0, int lvl = 0) {
    trav(u);
    int c = troid(u, csub[u]);
    done[c] = 1, cdepth[c] = lvl;
    cent[p].push_back(c);
    childs[c].push_back(c);
    nowh = lvl, nowc = c;
    for (auto &[v, t] : g[c])
      if (v ^ p && !done[v]) nowf = t, deep(v, c, t, 0, 0);
    for (auto &[v, t] : g[c])
      if (!done[v]) decomp(v, c, lvl + 1);
  }
  struct fenwick {
    int bit[sz];
    fenwick() { memset(bit, 0, sizeof bit); }
    void add(int x, int v) {
      x++;
      while (x <= n) {
        bit[x] += v;
        bit[x] >= mod ? bit[x] -= mod : 1;
        bit[x] < 0 ? bit[x] += mod : 1;
        x += x & -x;
      }
    }
    int quer(int x) {
      x++;
      int sum = 0;
      while (x) {
        sum += bit[x];
        sum >= mod ? sum -= mod : 1;
        sum < 0 ? sum += mod : 1;
        x -= x & -x;
      }
      return sum;
    }
  };
  int main() {
    int k;
    scanf("%d %d", &n, &k);
    for (int i = 1; i <= n; i++) scanf("%d", &a[i]);
    for (int i = 1; i < n; i++) {
      int u, v, t;
      scanf("%d %d %d", &u, &v, &t);
      g[u].emplace_back(v, t);
      g[v].emplace_back(u, t);
    }
    decomp(1);
    fenwick cnt[2], sum[2];
    long long ans = 0;
    for (int i = 1; i <= n; i++) {
      int h = cdepth[i];
      ans += a[i];
      ans %= mod;
      for (int j : cent[i]) {
        int f = fdis[h][j];
        for (int v : childs[j]) {
          long long dis = hdis[h][v] + a[i];
          int kcut = pdis[h][v];
          if (kcut > k) continue;
          kcut = k - kcut;
          ans += dis;
          ans += sum[f].quer(kcut);
          ans += sum[f ^ 1].quer(kcut - 1);
          ans += dis * cnt[f].quer(kcut) % mod;
          ans += dis * cnt[f ^ 1].quer(kcut - 1);
          ans %= mod;
        }
        for (int v : childs[j]) {
          int dis = hdis[h][v];
          int kcut = pdis[h][v];
          if (kcut > k) continue;
          sum[f].add(kcut, dis);
          cnt[f].add(kcut, +1);
        }
      }
      for (int j : cent[i]) {
        int f = fdis[h][j];
        for (int v : childs[j]) {
          int dis = hdis[h][v];
          int kcut = pdis[h][v];
          if (kcut > k) continue;
          sum[f].add(kcut, -dis);
          cnt[f].add(kcut, -1);
        }
      }
    }
    cout << ans;
  }
  