Problem p02710 (C++) — buggy submission, verdict: Runtime Error

/*"Success isn't permanent, failure isn't fatal,
                             it's the courage to continue that counts"*/
#include <bits/stdc++.h>
// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_policy.hpp>
// using namespace __gnu_pbds;
using namespace std;
// typedef tree<int, null_type, less<int>, rb_tree_tag,
// tree_order_statistics_node_update> ordered_set;
#define ll long long
#define scn(n) scanf("%d", &n)
#define lscn(n) scanf("%lld", &n)
#define lpri(n) printf("%lld", n)
#define pri(n) printf("%d", n)
#define pln() printf("\\n")
#define priln(n) printf("%d\\n", n)
#define lpriln(n) printf("%lld\\n", n)
#define rep(i, init, n) for (int i = init; i < n; i++)
#define pb push_back
#define mp make_pair
#define F first
#define S second
#define gcd __gcd
#define inf INT_MAX
#define ninf INT_MIN
const int mod = 1e9 + 7;
const int N = 2e5 + 4;

vector<int> v[N];

int col[N], sz[N];

ll ans[N], cur_size[N];

int dfs_util(int node, int par) {
  sz[node] = 1;

  for (int it : v[node])
    if (it != par)
      sz[node] += dfs_util(it, node);

  return sz[node];
}

void dfs(int node, int par) {
  ll n = cur_size[col[node]];

  ll contribution = 0, sum = 0;

  contribution =
      (n - sz[node] + 1) * (sz[node]); // paths form up and down or both

  for (int it : v[node])
    if (it != par) {
      // make the problem(temporarily) for counting new paths of col[node]
      // as a new problem for the subtree of it, that is make the size of graph
      // as sz[it],
      //  not the original 'n'.

      cur_size[col[node]] = sz[it];

      dfs(it, node);

      contribution += sum * sz[it]; // paths from subtree to subtree

      sum += sz[it];
    }

  ans[col[node]] += contribution;

  // now never take paths entering the subtree of node for col[node],
  // as we have considered all such paths.

  cur_size[col[node]] = n - sz[node];
}

int main() {
  int n;

  scn(n);

  rep(i, 1, n + 1) scn(col[i]), cur_size[i] = n;

  rep(i, 0, n - 1) {
    int a, b;

    scn(a);
    scn(b);

    v[a].pb(b);
    v[b].pb(a);
  }

  dfs_util(1, -1);

  dfs(1, -1);

  rep(i, 1, n + 1) lpriln(ans[i]);

  return 0;
}

Error: WA

Fixed submission:

/*"Success isn't permanent, failure isn't fatal,
                             it's the courage to continue that counts"*/
#include <bits/stdc++.h>
// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_policy.hpp>
// using namespace __gnu_pbds;
using namespace std;
// typedef tree<int, null_type, less<int>, rb_tree_tag,
// tree_order_statistics_node_update> ordered_set;
#define ll long long
#define scn(n) scanf("%d", &n)
#define lscn(n) scanf("%lld", &n)
#define lpri(n) printf("%lld", n)
#define pri(n) printf("%d", n)
#define pln() printf("\\n")
#define priln(n) printf("%d\\n", n)
#define lpriln(n) printf("%lld\\n", n)
#define rep(i, init, n) for (int i = init; i < n; i++)
#define pb push_back
#define mp make_pair
#define F first
#define S second
#define gcd __gcd
#define inf INT_MAX
#define ninf INT_MIN
const int mod = 1e9 + 7;
const int N = 2e3 + 4;

vector<pair<int, int>> v;

int n;

// dp[i][j] tells the best answer if we have processed till i-1 elements(and
// currently are at the ith element) and used the first j-1 places from the left
// to place some of the elements which we have taken now.

ll dp[N][N];

int main() {
  scn(n);

  rep(i, 0, n) {
    int val;

    scn(val);
    v.emplace_back(val, i);
  }

  // you will definetely try to give the largest element the most possible
  // movement.
  sort(v.rbegin(), v.rend());

  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++)
      dp[i][j] = (-1e18);

  dp[0][0] = 0;

  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      if (j <= i) {
        int right = j + (n - 1 - i);

        dp[i + 1][j + 1] =
            max(dp[i + 1][j + 1], dp[i][j] + v[i].F * abs(v[i].S - j));

        dp[i + 1][j] =
            max(dp[i + 1][j], dp[i][j] + v[i].F * abs(v[i].S - right));
      }
  }

  ll ans = 0;

  rep(i, 0, n + 1) ans = max(ans, dp[n][i]);

  lpri(ans);

  return 0;
}