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;
  template <typename T>
  struct fenwick {
    vector<T> a;
    int n;
    fenwick(int n) : n(n), a(n) {}
    void add(int x, T v) {
      assert(0 <= x && x < n);
      for (int i = x + 1; i <= n; i += i & -i) a[i - 1] += v;
    }
    T qry(int x) {
      T ret = 0;
      for (int i = min(x + 1, n); i > 0; i -= i & -i) ret += a[i - 1];
      return ret;
    }
  };
  const int mod = 1e9 + 7;
  long long inv(long long a) {
    long long u = 0, v = 1, b = mod;
    while (a) {
      long long t = b / a;
      b -= t * a, u -= t * v;
      swap(a, b);
      swap(u, v);
    }
    u %= mod;
    if (u < 0) u += mod;
    return u;
  }
  struct mint {
    int x;
    mint() : x(0) {}
    mint(int x) : x(x) {}
    mint(long long x) : x() {
      if (x >= mod || x < -mod) x %= mod;
      if (x < 0) x += mod;
      this->x = x;
    }
    mint &operator-=(const mint &rhs) {
      x -= rhs.x;
      if (x < 0) x += mod;
      return *this;
    }
    mint &operator+=(const mint &rhs) {
      x += rhs.x;
      if (x >= mod) x -= mod;
      return *this;
    }
    mint &operator*=(const mint &rhs) {
      x = (long long)x * rhs.x % mod;
      return *this;
    }
    mint &operator/=(const mint &rhs) {
      x = (long long)x * inv(rhs.x) % mod;
      return *this;
    }
    bool operator<(const mint &rhs) { return x < rhs.x; }
    bool operator==(const mint &rhs) { return x == rhs.x; }
    bool operator!=(const mint &rhs) { return x != rhs.x; }
    mint operator-() { return mint(x == 0 ? 0 : mod - x); }
    friend mint operator+(const mint &lhs, const mint &rhs) {
      mint ret(lhs);
      return ret += rhs;
    }
    friend mint operator-(const mint &lhs, const mint &rhs) {
      mint ret(lhs);
      return ret -= rhs;
    }
    friend mint operator*(const mint &lhs, const mint &rhs) {
      mint ret(lhs);
      return ret *= rhs;
    }
    friend mint operator/(const mint &lhs, const mint &rhs) {
      mint ret(lhs);
      return ret /= rhs;
    }
    friend ostream &operator<<(ostream &os, const mint &rhs) {
      os << rhs.x;
      return os;
    }
    friend istream &operator>>(istream &is, mint &rhs) {
      long long a;
      is >> a;
      rhs = a;
      return is;
    }
  };
  mint fpow(mint a, long long b) {
    mint r = 1;
    for (; b; b >>= 1) {
      if (b & 1) r = r * a;
      a = a * a;
    }
    return r;
  }
  int main() {
    int n, k;
    cin >> n >> k;
    ++k;
    vector<vector<pair<int, int>>> g(n);
    vector<int> a(n);
    for (int i = 0; i < n; ++i) cin >> a[i];
    for (int i = 1; i < n; ++i) {
      int u, v, w;
      cin >> u >> v >> w;
      --u, --v;
      g[u].push_back({v, w});
      g[v].push_back({u, w});
    }
    int rt = -1, res = n;
    vector<int> mx(n), sz(n), vis(n);
    function<void(int, int)> dfs = [&](int u, int fa) {
      sz[u] = 1, mx[u] = 0;
      for (auto [v, w] : g[u]) {
        if (v == fa || vis[v]) continue;
        dfs(v, u);
        sz[u] += sz[v];
        mx[u] = max(mx[u], sz[v]);
      }
      mx[u] = max(mx[u], res - sz[u]);
      if (rt == -1 || mx[u] < mx[rt]) rt = u;
    };
    mint ans = 0;
    for (int i = 0; i < n; ++i) ans += a[i];
    function<void(int)> dfs2 = [&](int u) {
      int pre = res;
      vis[u] = 1;
      fenwick<mint> fen(res);
      fenwick<int> fen2(res);
      map<pair<int, int>, mint> mp1;
      map<pair<int, int>, int> mp2;
      function<void(int, int, int, int, int, mint)> dfs3 =
          [&](int u, int fa, int pre, int cnt, int first, mint sum) {
            if (cnt > k) return;
            sum += a[u];
            mp1[{first, cnt}] += sum;
            mp2[{first, cnt}] += 1;
            fen.add(cnt, sum);
            fen2.add(cnt, 1);
            for (auto [v, w] : g[u]) {
              if (v == fa || vis[v]) continue;
              dfs3(v, u, w, cnt + (w != pre), first, sum);
            }
          };
      function<void(int, int, int, int, int, mint)> dfs4 = [&](int u, int fa,
                                                               int pre, int cnt,
                                                               int first,
                                                               mint sum) {
        if (cnt > k) return;
        sum += a[u];
        ans += fen.qry(k - cnt) + fen2.qry(k - cnt) * sum;
        mint tmp = fen.qry(k - cnt) + fen2.qry(k - cnt) * sum;
        if (mp1.count({first, k + 1 - cnt})) {
          ans += mp1[{first, k + 1 - cnt}] + mp2[{first, k + 1 - cnt}] * sum;
          mint tmp2 = mp1[{first, k + 1 - cnt}] + mp2[{first, k + 1 - cnt}] * sum;
        }
        for (auto [v, w] : g[u]) {
          if (v == fa || vis[v]) continue;
          dfs4(v, u, w, cnt + (w != pre), first, sum);
        }
      };
      fen.add(0, a[u]);
      fen2.add(0, 1);
      for (auto [v, w] : g[u]) {
        if (vis[v]) continue;
        dfs4(v, u, w, 1, w, 0);
        dfs3(v, u, w, 1, w, a[u]);
      }
      for (auto [v, w] : g[u]) {
        if (vis[v]) continue;
        res = sz[v] < sz[u] ? sz[v] : pre - sz[u];
        rt = -1;
        dfs(v, u);
        dfs2(rt);
      }
      res = pre;
    };
    rt = -1;
    dfs(0, -1);
    dfs2(rt);
    cout << ans << '\n';
  }
  