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 N = 200000 + 9;
  const int P = 1000000000 + 7;
  const int inv2 = (P + 1) / 2;
  inline int add(int a, int b) { return a + b < P ? a + b : a + b - P; }
  inline int sub(int a, int b) { return a - b < 0 ? a - b + P : a - b; }
  inline int mul(int a, int b) { return (long long)a * b % P; }
  int n, K, a[N], ans;
  vector<pair<int, int>> G[N];
  int rt, tot, msz, top, sz[N];
  bool vis[N];
  void dfs1(int u, int p) {
    int _max = 0;
    sz[u] = 1;
    for (auto [v, w] : G[u])
      if (!vis[v] && v != p) {
        dfs1(v, u);
        sz[u] += sz[v];
        _max = max(_max, sz[v]);
      }
    _max = max(_max, tot - sz[u]);
    if (_max < msz) {
      msz = _max;
      rt = u;
    }
  }
  vector<int> cnt[N], C[2], sum[N], S[2];
  void dfs3(int u, int p, int c, int k, int ans) {
    if (k > K) return;
    cnt[top][k] = add(cnt[top][k], 1);
    sum[top][k] = add(sum[top][k], ans);
    for (auto [v, w] : G[u])
      if (!vis[v] && v != p) {
        dfs3(v, u, w, k + (c != w), add(ans, a[v]));
      }
  }
  void dfs2(int u, int _tot) {
    tot = _tot;
    msz = tot + 1;
    dfs1(u, 0);
    vis[rt] = true;
    int _sz = sz[rt];
    C[0].resize(min(K + 1, _tot));
    fill(C[0].begin(), C[0].end(), 0);
    C[1].resize(min(K + 1, _tot));
    fill(C[1].begin(), C[1].end(), 0);
    S[0].resize(min(K + 1, _tot));
    fill(S[0].begin(), S[0].end(), 0);
    S[1].resize(min(K + 1, _tot));
    fill(S[1].begin(), S[1].end(), 0);
    for (auto [v, w] : G[rt])
      if (!vis[v]) {
        top = v;
        cnt[v].resize(min(K + 1, sz[v] < _sz ? sz[v] : _tot - _sz));
        fill(cnt[v].begin(), cnt[v].end(), 0);
        sum[v].resize(min(K + 1, sz[v] < _sz ? sz[v] : _tot - _sz));
        fill(sum[v].begin(), sum[v].end(), 0);
        dfs3(v, rt, w, 0, a[v]);
        for (int i = 0, ed = cnt[v].size(); i < ed; ++i) {
          C[w][i] = add(C[w][i], cnt[v][i]);
          S[w][i] = add(S[w][i], sum[v][i]);
        }
      }
    ans = add(ans, a[rt]);
    for (int i = 0, ed = C[0].size(); i < ed; ++i) {
      ans = add(ans, add(mul(a[rt], C[0][i]), S[0][i]));
      ans = add(ans, add(mul(a[rt], C[1][i]), S[1][i]));
    }
    for (int i = 1, ed = C[0].size(); i < ed; ++i) {
      C[0][i] = add(C[0][i], C[0][i - 1]);
      C[1][i] = add(C[1][i], C[1][i - 1]);
      S[0][i] = add(S[0][i], S[0][i - 1]);
      S[1][i] = add(S[1][i], S[1][i - 1]);
    }
    int tmp = 0;
    for (auto [v, w] : G[rt])
      if (!vis[v]) {
        for (int i = 1, ed = cnt[v].size(); i < ed; ++i) {
          cnt[v][i] = add(cnt[v][i], cnt[v][i - 1]);
          sum[v][i] = add(sum[v][i], sum[v][i - 1]);
        }
        for (int i = 0, ed = cnt[v].size(); i < ed; ++i) {
          C[w][i] = sub(C[w][i], cnt[v][i]);
          S[w][i] = sub(S[w][i], sum[v][i]);
        }
        int cntv_back = cnt[v].back();
        int sumv_back = sum[v].back();
        for (int i = cnt[v].size() - 1; i >= 1; --i) {
          cnt[v][i] = sub(cnt[v][i], cnt[v][i - 1]);
          sum[v][i] = sub(sum[v][i], sum[v][i - 1]);
        }
        for (int i = 0, ed = cnt[v].size(); i < ed; ++i) {
          int j = min(K - i, (int)C[0].size() - 1);
          tmp = add(tmp, mul(cnt[v][i], sub(S[w][j], j < ed ? 0 : sumv_back)));
          tmp = add(tmp, mul(sum[v][i], sub(C[w][j], j < ed ? 0 : cntv_back)));
          tmp = add(tmp, mul(a[rt], mul(cnt[v][i],
                                        sub(C[w][j], j < ed ? 0 : cntv_back))));
          if (i < K) {
            int j = min(K - i - 1, (int)C[0].size() - 1);
            tmp = add(tmp, mul(cnt[v][i], S[w ^ 1][j]));
            tmp = add(tmp, mul(sum[v][i], C[w ^ 1][j]));
            tmp = add(tmp, mul(a[rt], mul(cnt[v][i], C[w ^ 1][j])));
          }
        }
        for (int i = 1, ed = cnt[v].size(); i < ed; ++i) {
          cnt[v][i] = add(cnt[v][i], cnt[v][i - 1]);
          sum[v][i] = add(sum[v][i], sum[v][i - 1]);
        }
        for (int i = 0, ed = cnt[v].size(); i < ed; ++i) {
          C[w][i] = add(C[w][i], cnt[v][i]);
          S[w][i] = add(S[w][i], sum[v][i]);
        }
      }
    ans = add(ans, mul(tmp, inv2));
    for (auto [v, w] : G[rt])
      if (!vis[v]) {
        dfs2(v, sz[v] < _sz ? sz[v] : _tot - _sz);
      }
  }
  int main() {
    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, w;
      scanf("%d %d %d", &u, &v, &w);
      G[u].push_back({v, w});
      G[v].push_back({u, w});
    }
    dfs2(1, n);
    printf("%d\n", ans);
    return 0;
  }
  