Problem p02710 (C++) — buggy submission, verdict: Time Limit Exceeded

/**
 * code generated by JHelper
 * More info: https://github.com/AlexeyDmitriev/JHelper
 * @author Ido Kessler
 */

#include <iostream>

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cstdint>
#include <fstream>
#include <functional>
#include <iomanip>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <tuple>

using ll = long long int;
using vi = std::vector<int>;
using vvi = std::vector<vi>;
using vll = std::vector<ll>;
using vvl = std::vector<vll>;
using pi = std::pair<int, int>;
using vpi = std::vector<pi>;

#define all(x) (x).begin(), (x).end()
#define fori(i, n) for (int i = 0; i < n; i++)
#define rep(i, s, e) for (int i = (s); i <= (e); ++i)
#define repr(i, s, e) for (int i = (s); i >= (e); --i)
#define pb push_back

template <typename T> void mk_unique(T &l) {
  if (!std::is_sorted(all(l)))
    std::sort(all(l));
  auto it = std::unique(l.begin(), l.end());
  auto d = std::distance(l.begin(), it);
  l.resize(d);
}
template <typename T> T &smax(T &l, const T &r) {
  if (l < r)
    l = r;
  return l;
}
template <typename T> T &smin(T &l, const T &r) {
  if (l > r)
    l = r;
  return l;
}

constexpr int MOD = 1000000007;

namespace io {
class input_reader {
  void throw_error() {
    throw std::runtime_error("No more inputs while reading number");
  }
  template <typename T> inline void readIntegerValue(T &f) {
    bool positive = true;
    f = 0;
    char c;
    if (!in.get(c))
      throw_error();
    while (c < '0' || c > '9') {
      if (c == '-')
        positive = false;
      if (!in.get(c))
        throw_error();
    }
    while (c >= '0' && c <= '9') {
      f = (f << 3) + (f << 1) + (c & 15);
      if (!in.get(c))
        break;
    }
    if (!positive)
      f *= -1;
  }

public:
  std::istream &in;
  explicit input_reader(std::istream &in) : in(in) {}
  inline void read(int &f) { readIntegerValue(f); }
  inline void read(short &f) { readIntegerValue(f); }
  inline void read(long int &f) { readIntegerValue(f); }
  inline void read(long long int &f) { readIntegerValue(f); }
  inline void read(unsigned int &f) { readIntegerValue(f); }
  inline void read(unsigned short &f) { readIntegerValue(f); }
  inline void read(unsigned long int &f) { readIntegerValue(f); }
  inline void read(unsigned long long int &f) { readIntegerValue(f); }
  template <typename T> inline void read(std::vector<T> &f) {
    for (auto &i : f)
      read(i);
  }
  template <typename T, int N> inline void read(std::array<T, N> &res) {
    for (auto &i : res)
      read(i);
  }
  template <typename T1, typename T2> inline void read(std::pair<T1, T2> &p) {
    in.read(p.first, p.second);
  }
  template <typename T> inline void read(T &f) { in >> f; }
  template <typename T, typename... Args>
  inline void read(T &t, Args &...args) {
    read(t), read(args...);
  }
  template <typename T> inline void readArray(int n, T *res) {
    while (n--)
      read(*(res++));
  }
  template <typename T> inline input_reader &operator>>(T &t) {
    return read(t), *this;
  }
};
} // namespace io

namespace io {
class output_writer {
  std::ostream &out;
  template <typename T> inline void printIntegerHelper(T x) {
    if (x >= 10)
      printIntegerHelper(x / 10);
    out.put(x % 10 + 48);
  }
  template <typename T> inline void printInteger(T x) {
    if (x < 0)
      out.put('-'), x = -x;
    printIntegerHelper(x);
  }
  template <typename T> inline void printArray(T &ar) {
    bool first = true;
    for (auto &t : ar)
      print((first) ? first = false, "" : " ", t);
  }

public:
  explicit output_writer(std::ostream &out) : out(out) {}
  inline void flush() { out.flush(); }
  inline void print() {}
  inline void print(char x) { out.put(x); }
  inline void print(short x) { printInteger(x); }
  inline void print(int x) { printInteger(x); }
  inline void print(long int x) { printInteger(x); }
  inline void print(long long int x) { printInteger(x); }
  inline void print(unsigned short x) { printInteger(x); }
  inline void print(unsigned int x) { printInteger(x); }
  inline void print(unsigned long int x) { printInteger(x); }
  inline void print(unsigned long long int x) { printInteger(x); }
  template <typename T1> inline void print(const T1 &t1) { out << t1; }
  template <typename T1> inline void print(const std::vector<T1> &t1) {
    printArray(t1);
  }
  template <typename T1, int N> inline void print(const std::array<T1, N> &t1) {
    printArray(t1);
  }
  template <typename T1, int N, int M>
  inline void print(const std::array<std::array<T1, M>, N> &t1) {
    for (auto &t : t1)
      println(t);
  }
  template <typename T1>
  inline void print(const std::vector<std::vector<T1>> &t1) {
    bool first = true;
    for (auto &t : t1)
      print((first) ? first = false, "" : "\\n", t);
  }
  template <typename T, typename... Args>
  inline void print(const T &t, Args... args) {
    print(t), print(args...);
  }
  template <typename... Args> inline void println(Args... args) {
    print(args...), print('\\n');
  }
  template <typename T> output_writer &operator<<(const T &t) {
    return print(t), *this;
  }
};
} // namespace io

inline ll pow_mod(ll a, ll b, const ll MOD) {
  ll res = 1;
  for (ll i = 1; b; b ^= b & i, i <<= 1, a = (a * a) % MOD)
    if (b & i)
      res = (res * a) % MOD;
  return res;
}
inline ll mod_inv(const ll a, const ll MOD) { return pow_mod(a, MOD - 2, MOD); }
template <typename T> inline T gcd(T a, T b) {
  T t;
  while (a != 0)
    t = a, a = b % a, b = t;
  return b;
}
template <typename T> T lcm(const T &a, const T &b) {
  assert(a != 0 && b != 0);
  return a / gcd(a, b) * b;
}

namespace math {
class mod_int {
  int v;

public:
  mod_int(int v) : v(v % MOD) {}
  mod_int(unsigned int v) : v(v % MOD) {}
  mod_int(long long int v) : v((int)(v % MOD)) {}
  mod_int(unsigned long long int v) : v((int)(v % MOD)) {}
  mod_int(long int v) : v((int)(v % MOD)) {}
  mod_int(unsigned long int v) : v((int)(v % MOD)) {}
  mod_int() = default;
  mod_int(const mod_int &o) = default;
  mod_int(mod_int &&o) = default;
  inline mod_int &operator=(const mod_int &o) = default;
  inline mod_int &operator=(mod_int &&o) = default;

  inline mod_int &operator++() {
    if (v == MOD - 1)
      v = 0;
    else
      v++;
    return *this;
  }
  inline mod_int &operator--() {
    if (v == 0)
      v = MOD - 1;
    else
      v--;
    return *this;
  }
  inline mod_int operator++(int) {
    mod_int t = *this;
    ++(*this);
    return t;
  }
  inline mod_int operator--(int) {
    mod_int t = *this;
    --(*this);
    return t;
  }

  // inline bool friend operator<(const mod_int&lhs,const mod_int&rhs) { return
  // lhs.v  <  rhs.v; } inline bool friend operator>(const mod_int&lhs,const
  // mod_int&rhs) { return lhs.v  <  rhs.v; } inline bool friend
  // operator<=(const mod_int&lhs,const mod_int&rhs) { return lhs.v <= rhs.v; }
  // inline bool friend operator>=(const mod_int&lhs,const mod_int&rhs) { return
  // lhs.v <= rhs.v; }
  inline bool friend operator==(const mod_int &lhs, const mod_int &rhs) {
    return lhs.v == rhs.v;
  }
  inline bool friend operator!=(const mod_int &lhs, const mod_int &rhs) {
    return lhs.v != rhs.v;
  }
  // template<typename T> inline bool friend operator<(const mod_int&lhs,const
  // T&rhs) { return  lhs <  mod_int(rhs); } template<typename T> inline bool
  // friend operator>(const mod_int&lhs,const T&rhs) { return  lhs >
  // mod_int(rhs); } template<typename T> inline bool friend operator<=(const
  // mod_int&lhs,const T&rhs) { return lhs <= mod_int(rhs); }
  // template<typename T> inline bool friend operator>=(const mod_int&lhs,const
  // T&rhs) { return lhs >= mod_int(rhs); }
  template <typename T>
  inline bool friend operator==(const mod_int &lhs, const T &rhs) {
    return lhs == mod_int(rhs);
  }
  template <typename T>
  inline bool friend operator!=(const mod_int &lhs, const T &rhs) {
    return lhs != mod_int(rhs);
  }

  inline mod_int friend operator-(const mod_int &lhs, const mod_int &rhs) {
    int v = lhs.v - rhs.v;
    if (v < 0)
      v += MOD;
    return {v};
  }
  inline mod_int friend operator+(const mod_int &lhs, const mod_int &rhs) {
    int v = lhs.v + rhs.v;
    if (v >= MOD)
      v -= MOD;
    return {v};
  }
  inline mod_int friend operator*(const mod_int &lhs, const mod_int &rhs) {
    return {(int)((((long long int)lhs.v) * rhs.v) % MOD)};
  }
  inline mod_int friend operator/(const mod_int &lhs, const mod_int &rhs) {
    return {(int)((((long long int)lhs.v) * rhs.inv().v) % MOD)};
  }
  inline mod_int friend operator%(const mod_int &lhs, const mod_int &rhs) {
    return {lhs.v % rhs.v};
  }

  template <typename T>
  inline mod_int friend operator-(const mod_int &lhs, const T &rhs) {
    return lhs - mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator+(const mod_int &lhs, const T &rhs) {
    return lhs + mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator*(const mod_int &lhs, const T &rhs) {
    return lhs * mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator/(const mod_int &lhs, const T &rhs) {
    return lhs / mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator%(const mod_int &lhs, const T &rhs) {
    return {lhs.v % rhs};
  }

  template <typename T> inline mod_int &operator-=(const T &rhs) {
    return *this = *this - rhs;
  }
  template <typename T> inline mod_int &operator+=(const T &rhs) {
    return *this = *this + rhs;
  }
  template <typename T> inline mod_int &operator*=(const T &rhs) {
    return *this = *this * rhs;
  }
  template <typename T> inline mod_int &operator/=(const T &rhs) {
    return *this = *this / rhs;
  }
  template <typename T> inline mod_int &operator%=(const T &rhs) {
    return *this = *this % rhs;
  }

  inline mod_int inv() const { return pow(MOD - 2); }
  inline mod_int pow(long long int pw) const {
    mod_int res = 1, val = v;
    while (pw) {
      if (pw & 1)
        res *= val;
      val *= val;
      pw >>= 1;
    }
    return res;
  }

  inline operator int() const { return v; }

  inline friend std::istream &operator>>(std::istream &in, mod_int &rhs) {
    in >> rhs.v;
    rhs.v %= MOD;
    if (rhs.v < 0)
      rhs.v += MOD;
    return in;
  }
  inline friend std::ostream &operator<<(std::ostream &out, mod_int &rhs) {
    out << rhs.v;
    return out;
  }
};
} // namespace math

namespace math {
mod_int factorial(int n) {
  static std::vector<mod_int> mem;
  static mod_int last;
  if (mem.empty())
    last = {1}, mem.push_back(last);
  while (mem.size() <= n)
    mem.push_back(mem.back() * last++);
  return mem[n];
}

mod_int n_choose_k(int n, int k) {
  mod_int res = factorial(n);
  mod_int res2 = (factorial(k) * factorial(n - k));
  return res / res2;
}
} // namespace math
// #include "points_query/point_query.hpp"
// #include "numeric/numeric.hpp"
// #include "hashes/hashes.hpp"
// #include "graphs/graphs.hpp"
// #include "geometry/geometry.hpp"
// #include "extension_operations/extensions.hpp"
// #include "ranges/ranges.hpp"
using mi = math::mod_int;

using namespace std;

class FPathPassI {
  class Solver {
  public:
    io::input_reader in;
    io::output_writer out;

    Solver(std::istream &in, std::ostream &out) : in(in), out(out) {}

    void print_case_number() {
      static int t = 0;
      //            out <<"Case #" << ++t << ": ";
    }

    vi q;
    vvi e;
    vi c;
    vector<ll> ans;
    vector<bool> was_center;
    int n;
    vi max_subtree;
    vi no_children;

    int calc_num_of_children(int i, int p) {
      int no_children_ = 1;
      for (int j : e[i]) {
        if (j != p && !was_center[j]) {
          no_children_ += calc_num_of_children(j, i);
        }
      }
      no_children[i] = no_children_;
      return no_children_;
    }

    void calc_centroid_and_max_sub_tree(int i, int p, int n_cur_tree,
                                        int &centroid_ind) {
      int max_subtree_ = n_cur_tree - no_children[i];
      for (int j : e[i]) {
        if (j != p && !was_center[j]) {
          calc_centroid_and_max_sub_tree(j, i, n_cur_tree, centroid_ind);
          smax(max_subtree_, no_children[j]);
        }
      }
      if (max_subtree[centroid_ind] > max_subtree_) {
        centroid_ind = i;
        max_subtree[i] = max_subtree_;
      }
    }

    int get_centroid(int i) {
      calc_num_of_children(i, -1);
      int centroid = i;
      max_subtree[i] = no_children[i];
      calc_centroid_and_max_sub_tree(i, -1, no_children[i], centroid);
      return centroid;
    }

    ll dfs(int i, int p) {
      ll count_paths = 1;
      bool is_first = !cur_path_colors[c[i]];
      if (is_first)
        cur_path_colors[c[i]] = true;
      for (int j : e[i]) {
        if (p != j && !was_center[j]) {
          count_paths += dfs(j, i);
        }
      }
      if (is_first)
        sketch1[c[i]] += count_paths, cur_path_colors[c[i]] = false;
      return count_paths;
    }

    vector<bool> cur_path_colors;

    void calculate_all_paths_starting_from(int i, int centroid_color) {
      if (cur_path_colors.size() != n)
        cur_path_colors.resize(n);
      sketch1.clear();
      cur_path_colors[centroid_color] = true;
      sketch1[centroid_color] = dfs(i, -1);
      cur_path_colors[centroid_color] = false;
    }

    template <class T> class FastMap {
    public:
      vi indices;

    private:
      vector<pair<T, int>> data_version;
      int version = 0;

    public:
      void clear() {
        version++;
        indices.clear();
      }

      void resize(int sz) { data_version.resize(sz); }

      T &operator[](int i) {
        auto &x = data_version[i];
        if (x.second != version) {
          x = {0, version};
          indices.pb(i);
        }
        return x.first;
      }
    };

    FastMap<int> sketch;
    FastMap<int> sketch1;

    void calculate_all_paths_passing_thoruth(int i) {
      sketch.clear();
      calc_num_of_children(i, -1);
      was_center[i] = true;
      sketch[c[i]] = 1;
      ll total_left = no_children[i];
      ll total_passed = 0;
      ans[c[i]] += 1 * total_left;
      total_left--;
      total_passed++;
      for (int j : e[i]) {
        if (was_center[j])
          continue;
        calculate_all_paths_starting_from(j, c[i]);
        total_left -= no_children[j];
        for (auto k : sketch1.indices) {
          ans[k] += 1ll * sketch1[k] * total_left;
          ans[k] += 1ll * (total_passed - (sketch[k])) * sketch1[k];
          sketch[k] += sketch1[k];
        }
        total_passed += no_children[j];
      }
    }

    void solve_rec(int i) {
      if (was_center[i])
        return;
      i = get_centroid(i);
      calculate_all_paths_passing_thoruth(i);
      for (int j : e[i])
        if (!was_center[j])
          q.pb(j);
    }

    void solve() {
      print_case_number();
      // <code here>
      in >> n;
      e.resize(n);
      c.resize(n);
      ans.resize(n);
      was_center.resize(n);
      no_children.resize(n);
      max_subtree.resize(n);
      sketch.resize(n);
      sketch1.resize(n);
      in >> c;
      fori(i, n) c[i]--;
      fori(i, n - 1) {
        int u, v;
        in.read(u, v);
        u--, v--;
        e[u].pb(v), e[v].pb(u);
      }
      q.clear();
      q.pb(0);
      while (!q.empty()) {
        int qq = q.back();
        q.pop_back();
        solve_rec(qq);
      }

      // </code here>
      out.println(ans);
    }
  };

public:
  void solve(std::istream &in, std::ostream &out) { Solver(in, out).solve(); }
};

int main() {
  std::ios_base::sync_with_stdio(false), cin.tie(0);
  FPathPassI solver;
  std::istream &in(std::cin);
  std::ostream &out(std::cout);
  solver.solve(in, out);
  return 0;
}


Error: TLE

Fixed submission:

/**
 * code generated by JHelper
 * More info: https://github.com/AlexeyDmitriev/JHelper
 * @author Ido Kessler
 */

#include <iostream>

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cstdint>
#include <fstream>
#include <functional>
#include <iomanip>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <tuple>

using ll = long long int;
using vi = std::vector<int>;
using vvi = std::vector<vi>;
using vll = std::vector<ll>;
using vvl = std::vector<vll>;
using pi = std::pair<int, int>;
using vpi = std::vector<pi>;

#define all(x) (x).begin(), (x).end()
#define fori(i, n) for (int i = 0; i < n; i++)
#define rep(i, s, e) for (int i = (s); i <= (e); ++i)
#define repr(i, s, e) for (int i = (s); i >= (e); --i)
#define pb push_back

template <typename T> void mk_unique(T &l) {
  if (!std::is_sorted(all(l)))
    std::sort(all(l));
  auto it = std::unique(l.begin(), l.end());
  auto d = std::distance(l.begin(), it);
  l.resize(d);
}
template <typename T> T &smax(T &l, const T &r) {
  if (l < r)
    l = r;
  return l;
}
template <typename T> T &smin(T &l, const T &r) {
  if (l > r)
    l = r;
  return l;
}

constexpr int MOD = 1000000007;

namespace io {
class input_reader {
  void throw_error() {
    throw std::runtime_error("No more inputs while reading number");
  }
  template <typename T> inline void readIntegerValue(T &f) {
    bool positive = true;
    f = 0;
    char c;
    if (!in.get(c))
      throw_error();
    while (c < '0' || c > '9') {
      if (c == '-')
        positive = false;
      if (!in.get(c))
        throw_error();
    }
    while (c >= '0' && c <= '9') {
      f = (f << 3) + (f << 1) + (c & 15);
      if (!in.get(c))
        break;
    }
    if (!positive)
      f *= -1;
  }

public:
  std::istream &in;
  explicit input_reader(std::istream &in) : in(in) {}
  inline void read(int &f) { readIntegerValue(f); }
  inline void read(short &f) { readIntegerValue(f); }
  inline void read(long int &f) { readIntegerValue(f); }
  inline void read(long long int &f) { readIntegerValue(f); }
  inline void read(unsigned int &f) { readIntegerValue(f); }
  inline void read(unsigned short &f) { readIntegerValue(f); }
  inline void read(unsigned long int &f) { readIntegerValue(f); }
  inline void read(unsigned long long int &f) { readIntegerValue(f); }
  template <typename T> inline void read(std::vector<T> &f) {
    for (auto &i : f)
      read(i);
  }
  template <typename T, int N> inline void read(std::array<T, N> &res) {
    for (auto &i : res)
      read(i);
  }
  template <typename T1, typename T2> inline void read(std::pair<T1, T2> &p) {
    in.read(p.first, p.second);
  }
  template <typename T> inline void read(T &f) { in >> f; }
  template <typename T, typename... Args>
  inline void read(T &t, Args &...args) {
    read(t), read(args...);
  }
  template <typename T> inline void readArray(int n, T *res) {
    while (n--)
      read(*(res++));
  }
  template <typename T> inline input_reader &operator>>(T &t) {
    return read(t), *this;
  }
};
} // namespace io

namespace io {
class output_writer {
  std::ostream &out;
  template <typename T> inline void printIntegerHelper(T x) {
    if (x >= 10)
      printIntegerHelper(x / 10);
    out.put(x % 10 + 48);
  }
  template <typename T> inline void printInteger(T x) {
    if (x < 0)
      out.put('-'), x = -x;
    printIntegerHelper(x);
  }
  template <typename T> inline void printArray(T &ar) {
    bool first = true;
    for (auto &t : ar)
      print((first) ? first = false, "" : " ", t);
  }

public:
  explicit output_writer(std::ostream &out) : out(out) {}
  inline void flush() { out.flush(); }
  inline void print() {}
  inline void print(char x) { out.put(x); }
  inline void print(short x) { printInteger(x); }
  inline void print(int x) { printInteger(x); }
  inline void print(long int x) { printInteger(x); }
  inline void print(long long int x) { printInteger(x); }
  inline void print(unsigned short x) { printInteger(x); }
  inline void print(unsigned int x) { printInteger(x); }
  inline void print(unsigned long int x) { printInteger(x); }
  inline void print(unsigned long long int x) { printInteger(x); }
  template <typename T1> inline void print(const T1 &t1) { out << t1; }
  template <typename T1> inline void print(const std::vector<T1> &t1) {
    printArray(t1);
  }
  template <typename T1, int N> inline void print(const std::array<T1, N> &t1) {
    printArray(t1);
  }
  template <typename T1, int N, int M>
  inline void print(const std::array<std::array<T1, M>, N> &t1) {
    for (auto &t : t1)
      println(t);
  }
  template <typename T1>
  inline void print(const std::vector<std::vector<T1>> &t1) {
    bool first = true;
    for (auto &t : t1)
      print((first) ? first = false, "" : "\\n", t);
  }
  template <typename T, typename... Args>
  inline void print(const T &t, Args... args) {
    print(t), print(args...);
  }
  template <typename... Args> inline void println(Args... args) {
    print(args...), print('\\n');
  }
  template <typename T> output_writer &operator<<(const T &t) {
    return print(t), *this;
  }
};
} // namespace io

inline ll pow_mod(ll a, ll b, const ll MOD) {
  ll res = 1;
  for (ll i = 1; b; b ^= b & i, i <<= 1, a = (a * a) % MOD)
    if (b & i)
      res = (res * a) % MOD;
  return res;
}
inline ll mod_inv(const ll a, const ll MOD) { return pow_mod(a, MOD - 2, MOD); }
template <typename T> inline T gcd(T a, T b) {
  T t;
  while (a != 0)
    t = a, a = b % a, b = t;
  return b;
}
template <typename T> T lcm(const T &a, const T &b) {
  assert(a != 0 && b != 0);
  return a / gcd(a, b) * b;
}

namespace math {
class mod_int {
  int v;

public:
  mod_int(int v) : v(v % MOD) {}
  mod_int(unsigned int v) : v(v % MOD) {}
  mod_int(long long int v) : v((int)(v % MOD)) {}
  mod_int(unsigned long long int v) : v((int)(v % MOD)) {}
  mod_int(long int v) : v((int)(v % MOD)) {}
  mod_int(unsigned long int v) : v((int)(v % MOD)) {}
  mod_int() = default;
  mod_int(const mod_int &o) = default;
  mod_int(mod_int &&o) = default;
  inline mod_int &operator=(const mod_int &o) = default;
  inline mod_int &operator=(mod_int &&o) = default;

  inline mod_int &operator++() {
    if (v == MOD - 1)
      v = 0;
    else
      v++;
    return *this;
  }
  inline mod_int &operator--() {
    if (v == 0)
      v = MOD - 1;
    else
      v--;
    return *this;
  }
  inline mod_int operator++(int) {
    mod_int t = *this;
    ++(*this);
    return t;
  }
  inline mod_int operator--(int) {
    mod_int t = *this;
    --(*this);
    return t;
  }

  // inline bool friend operator<(const mod_int&lhs,const mod_int&rhs) { return
  // lhs.v  <  rhs.v; } inline bool friend operator>(const mod_int&lhs,const
  // mod_int&rhs) { return lhs.v  <  rhs.v; } inline bool friend
  // operator<=(const mod_int&lhs,const mod_int&rhs) { return lhs.v <= rhs.v; }
  // inline bool friend operator>=(const mod_int&lhs,const mod_int&rhs) { return
  // lhs.v <= rhs.v; }
  inline bool friend operator==(const mod_int &lhs, const mod_int &rhs) {
    return lhs.v == rhs.v;
  }
  inline bool friend operator!=(const mod_int &lhs, const mod_int &rhs) {
    return lhs.v != rhs.v;
  }
  // template<typename T> inline bool friend operator<(const mod_int&lhs,const
  // T&rhs) { return  lhs <  mod_int(rhs); } template<typename T> inline bool
  // friend operator>(const mod_int&lhs,const T&rhs) { return  lhs >
  // mod_int(rhs); } template<typename T> inline bool friend operator<=(const
  // mod_int&lhs,const T&rhs) { return lhs <= mod_int(rhs); }
  // template<typename T> inline bool friend operator>=(const mod_int&lhs,const
  // T&rhs) { return lhs >= mod_int(rhs); }
  template <typename T>
  inline bool friend operator==(const mod_int &lhs, const T &rhs) {
    return lhs == mod_int(rhs);
  }
  template <typename T>
  inline bool friend operator!=(const mod_int &lhs, const T &rhs) {
    return lhs != mod_int(rhs);
  }

  inline mod_int friend operator-(const mod_int &lhs, const mod_int &rhs) {
    int v = lhs.v - rhs.v;
    if (v < 0)
      v += MOD;
    return {v};
  }
  inline mod_int friend operator+(const mod_int &lhs, const mod_int &rhs) {
    int v = lhs.v + rhs.v;
    if (v >= MOD)
      v -= MOD;
    return {v};
  }
  inline mod_int friend operator*(const mod_int &lhs, const mod_int &rhs) {
    return {(int)((((long long int)lhs.v) * rhs.v) % MOD)};
  }
  inline mod_int friend operator/(const mod_int &lhs, const mod_int &rhs) {
    return {(int)((((long long int)lhs.v) * rhs.inv().v) % MOD)};
  }
  inline mod_int friend operator%(const mod_int &lhs, const mod_int &rhs) {
    return {lhs.v % rhs.v};
  }

  template <typename T>
  inline mod_int friend operator-(const mod_int &lhs, const T &rhs) {
    return lhs - mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator+(const mod_int &lhs, const T &rhs) {
    return lhs + mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator*(const mod_int &lhs, const T &rhs) {
    return lhs * mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator/(const mod_int &lhs, const T &rhs) {
    return lhs / mod_int(rhs);
  }
  template <typename T>
  inline mod_int friend operator%(const mod_int &lhs, const T &rhs) {
    return {lhs.v % rhs};
  }

  template <typename T> inline mod_int &operator-=(const T &rhs) {
    return *this = *this - rhs;
  }
  template <typename T> inline mod_int &operator+=(const T &rhs) {
    return *this = *this + rhs;
  }
  template <typename T> inline mod_int &operator*=(const T &rhs) {
    return *this = *this * rhs;
  }
  template <typename T> inline mod_int &operator/=(const T &rhs) {
    return *this = *this / rhs;
  }
  template <typename T> inline mod_int &operator%=(const T &rhs) {
    return *this = *this % rhs;
  }

  inline mod_int inv() const { return pow(MOD - 2); }
  inline mod_int pow(long long int pw) const {
    mod_int res = 1, val = v;
    while (pw) {
      if (pw & 1)
        res *= val;
      val *= val;
      pw >>= 1;
    }
    return res;
  }

  inline operator int() const { return v; }

  inline friend std::istream &operator>>(std::istream &in, mod_int &rhs) {
    in >> rhs.v;
    rhs.v %= MOD;
    if (rhs.v < 0)
      rhs.v += MOD;
    return in;
  }
  inline friend std::ostream &operator<<(std::ostream &out, mod_int &rhs) {
    out << rhs.v;
    return out;
  }
};
} // namespace math

namespace math {
mod_int factorial(int n) {
  static std::vector<mod_int> mem;
  static mod_int last;
  if (mem.empty())
    last = {1}, mem.push_back(last);
  while (mem.size() <= n)
    mem.push_back(mem.back() * last++);
  return mem[n];
}

mod_int n_choose_k(int n, int k) {
  mod_int res = factorial(n);
  mod_int res2 = (factorial(k) * factorial(n - k));
  return res / res2;
}
} // namespace math
// #include "points_query/point_query.hpp"
// #include "numeric/numeric.hpp"
// #include "hashes/hashes.hpp"
// #include "graphs/graphs.hpp"
// #include "geometry/geometry.hpp"
// #include "extension_operations/extensions.hpp"
// #include "ranges/ranges.hpp"
using mi = math::mod_int;

using namespace std;

class FPathPassI {
  class Solver {
  public:
    io::input_reader in;
    io::output_writer out;

    Solver(std::istream &in, std::ostream &out) : in(in), out(out) {}

    void print_case_number() {
      static int t = 0;
      //            out <<"Case #" << ++t << ": ";
    }

    vi q;
    vvi e;
    vi c;
    vector<ll> ans;
    vector<bool> was_center;
    int n;
    vi max_subtree;
    vi no_children;

    int calc_num_of_children(int i, int p) {
      int no_children_ = 1;
      for (int j : e[i]) {
        if (j != p && !was_center[j]) {
          no_children_ += calc_num_of_children(j, i);
        }
      }
      no_children[i] = no_children_;
      return no_children_;
    }

    void calc_centroid_and_max_sub_tree(int i, int p, int n_cur_tree,
                                        int &centroid_ind) {
      int max_subtree_ = n_cur_tree - no_children[i];
      for (int j : e[i]) {
        if (j != p && !was_center[j]) {
          calc_centroid_and_max_sub_tree(j, i, n_cur_tree, centroid_ind);
          smax(max_subtree_, no_children[j]);
        }
      }
      if (max_subtree[centroid_ind] > max_subtree_) {
        centroid_ind = i;
        max_subtree[i] = max_subtree_;
      }
    }

    int get_centroid(int i) {
      calc_num_of_children(i, -1);
      int centroid = i;
      max_subtree[i] = no_children[i];
      calc_centroid_and_max_sub_tree(i, -1, no_children[i], centroid);
      return centroid;
    }

    ll dfs(int i, int p, map<int, ll> &res, set<int> cur_path) {
      ll count_paths = 1;
      bool is_first = cur_path.find(c[i]) == cur_path.end();
      if (is_first)
        cur_path.insert(c[i]);
      for (int j : e[i]) {
        if (p != j && !was_center[j]) {
          count_paths += dfs(j, i, res, cur_path);
        }
      }
      if (is_first)
        res[c[i]] += count_paths;
      return count_paths;
    }

    map<int, ll> calculate_all_paths_starting_from(int i, int centroid_color) {
      map<int, ll> res;
      set<int> cur_path;
      cur_path.insert(centroid_color);
      res[centroid_color] = dfs(i, -1, res, cur_path);
      return res;
    }

    void calculate_all_paths_passing_thoruth(int i) {
      calc_num_of_children(i, -1);
      was_center[i] = true;

      map<int, ll> sketch;
      sketch[c[i]] = 1;
      ll total_left = no_children[i];
      ll total_passed = 0;
      ans[c[i]] += 1 * total_left;
      total_left--;
      total_passed++;
      for (int j : e[i]) {
        if (was_center[j])
          continue;
        auto m = calculate_all_paths_starting_from(j, c[i]);
        total_left -= no_children[j];
        for (auto k : m) {
          ans[k.first] += k.second * total_left;
          ans[k.first] += (total_passed - (sketch[k.first])) * k.second;
          sketch[k.first] += k.second;
        }
        total_passed += no_children[j];
      }
    }

    void solve_rec(int i) {
      if (was_center[i])
        return;
      i = get_centroid(i);
      calculate_all_paths_passing_thoruth(i);
      for (int j : e[i])
        if (!was_center[j])
          q.pb(j);
    }

    void solve() {
      print_case_number();
      // <code here>
      in >> n;
      e.resize(n);
      c.resize(n);
      ans.resize(n);
      was_center.resize(n);
      no_children.resize(n);
      max_subtree.resize(n);

      in >> c;
      fori(i, n) c[i]--;
      fori(i, n - 1) {
        int u, v;
        in.read(u, v);
        u--, v--;
        e[u].pb(v), e[v].pb(u);
      }
      q.clear();
      q.pb(0);
      while (!q.empty()) {
        int qq = q.back();
        q.pop_back();
        solve_rec(qq);
      }

      // </code here>
      out.println(ans);
    }
  };

public:
  void solve(std::istream &in, std::ostream &out) { Solver(in, out).solve(); }
};

int main() {
  std::ios_base::sync_with_stdio(false), cin.tie(0);
  FPathPassI solver;
  std::istream &in(std::cin);
  std::ostream &out(std::cout);
  solver.solve(in, out);
  return 0;
}