prithivMLmods/Coder-Stat huggingface dataset:
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <iostream>
#include <utility>
using namespace std;

int main() {
  pair<long long, int> order[2020];
  int n;
  cin >> n;
  for (int i = 0; i < n; i++) {
    long long temp;
    cin >> temp;
    order[i] = make_pair(temp, i + 1);
  }
  sort(order, order + n);
  reverse(order, order + n);
  long long dp[1000][1000];
  for (int x = 0; x <= n; x++) {
    for (int y = 0; y <= n; y++) {
      dp[x][y] = -1;
    }
  }
  dp[0][0] = 0;
  for (int i = 0; i < n; i++) {
    long long k = order[i].first;
    int bef = order[i].second;
    for (int x = 0; x <= i; x++) {
      int y = i - x;
      dp[x + 1][y] = max(dp[x + 1][y], dp[x][y] + abs(bef - (x + 1)) * k);
      dp[x][y + 1] = max(dp[x][y + 1], dp[x][y] + abs(bef - (n - y)) * k);
    }
  }
  long long sum = 0;
  for (int i = 0; i <= n; i++) {
    sum = max(sum, dp[i][n - i]);
  }
  cout << sum << endl;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <iostream>
#include <utility>
using namespace std;

int main() {
  pair<long long, int> order[2020];
  int n;
  cin >> n;
  for (int i = 0; i < n; i++) {
    long long temp;
    cin >> temp;
    order[i] = make_pair(temp, i + 1);
  }
  sort(order, order + n);
  reverse(order, order + n);
  long long dp[2020][2020];
  for (int x = 0; x <= n; x++) {
    for (int y = 0; y <= n; y++) {
      dp[x][y] = -1;
    }
  }
  dp[0][0] = 0;
  for (int i = 0; i < n; i++) {
    long long k = order[i].first;
    int bef = order[i].second;
    for (int x = 0; x <= i; x++) {
      int y = i - x;
      dp[x + 1][y] = max(dp[x + 1][y], dp[x][y] + abs(bef - (x + 1)) * k);
      dp[x][y + 1] = max(dp[x][y + 1], dp[x][y] + abs(bef - (n - y)) * k);
    }
  }
  long long sum = 0;
  for (int i = 0; i <= n; i++) {
    sum = max(sum, dp[i][n - i]);
  }
  cout << sum << endl;
  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
const int mod = 1e9 + 7;
typedef long long ll;
const int N = 2e3 + 5;

int n;
std::vector<std::pair<ll, int>> v;
ll dp[N][N];
ll solve(int l, int r) {
  if (dp[l][r] >= 0)
    return dp[l][r];
  if (r < l)
    return dp[l][r] = 0;
  int i = n - r + l; // 当前正在放第i个
  return std::max(v[i].first * std::abs(v[i].second - l) + solve(l + 1, r),
                  v[i].first * std::abs(v[i].second - r) + solve(l, r - 1));
}
int main() {
  // freopen("a.in","r",stdin);
  scanf("%d", &n);
  v.push_back(std::make_pair(0, 0));
  for (int i = 1; i <= n; i++) {
    ll j;
    scanf("%lld", &j);
    v.push_back(std::make_pair(j, i));
  }
  std::sort(v.rbegin(), v.rend() - 1);
  memset(dp, -1, sizeof(dp));
  printf("%lld\n", solve(1, n));
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
const int mod = 1e9 + 7;
typedef long long ll;
const int N = 2e3 + 5;

int n;
std::vector<std::pair<ll, int>> v;
ll dp[N][N];
ll solve(int l, int r) {
  if (dp[l][r] >= 0)
    return dp[l][r];
  if (r < l)
    return dp[l][r] = 0;
  int i = n - r + l; // 当前正在放第i个
  return dp[l][r] =
             std::max(v[i].first * std::abs(v[i].second - l) + solve(l + 1, r),
                      v[i].first * std::abs(v[i].second - r) + solve(l, r - 1));
}
int main() {
  // freopen("a.in","r",stdin);
  scanf("%d", &n);
  v.push_back(std::make_pair(0, 0));
  for (int i = 1; i <= n; i++) {
    ll j;
    scanf("%lld", &j);
    v.push_back(std::make_pair(j, i));
  }
  std::sort(v.rbegin(), v.rend() - 1);
  memset(dp, -1, sizeof(dp));
  printf("%lld\n", solve(1, n));
  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

typedef long long ll;
typedef pair<ll, int> P;
int n;
vector<P> A;
ll dp[2005][2005];
ll nex[2005][2005];
ll ans;
int main() {

  cin >> n;

  for (int i = 0; i < n; i++) {
    ll a;
    cin >> a;

    A.push_back(P(a, i));
  }

  sort(A.begin(), A.end(), greater<P>());

  // for(int i = 0; i < n; i++) cout << A[i].first <<" "<<A[i].second<< endl;

  for (ll i = 0; i < n; i++) {
    for (ll j = 0; j <= i + 1; j++) {
      ll k = i - j + 1;
      if (j == 0) {
        nex[j][k] = dp[j][k - 1] + abs(n - k - A[i].second) * A[i].first;
        continue;
      }
      if (k == 0) {
        nex[j][k] = dp[j - 1][k] + abs(A[i].second - j + 1) * A[i].first;
        continue;
      }
      nex[j][k] = max(dp[j - 1][k] + abs(A[i].second - j + 1) * A[i].first,
                      dp[j][k - 1] + abs(n - k - A[i].second) * A[i].first);
    }
    swap(nex, dp);
    /* cout << i << endl;
     for(int ii = 0; ii <= i+1; ii++) {
       int j = i - ii + 1;
    cout <<"dp["<<ii<<"]["<<j<<"] = "<< dp[ii][j] << endl;
    }*/
  }

  for (int i = 0; i <= n; i++) {
    int j = n - i;
    ans = max(ans, dp[i][j]);
    // cout <<"dp["<<i<<"]["<<j<<"] = "<< dp[i][j] << endl;
  }
  cout << ans << endl;
}


Fixed submission:
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

typedef long long ll;
typedef pair<ll, int> P;
int n;
vector<P> A;
ll dp[2005][2005];
ll nex[2005][2005];
ll ans;
int main() {

  cin >> n;

  for (int i = 0; i < n; i++) {
    ll a;
    cin >> a;

    A.push_back(P(a, i));
  }

  sort(A.begin(), A.end(), greater<P>());

  // for(int i = 0; i < n; i++) cout << A[i].first <<" "<<A[i].second<< endl;

  for (ll i = 0; i < n; i++) {
    for (ll j = 0; j <= i + 1; j++) {
      ll k = i - j + 1;
      if (j == 0) {
        nex[j][k] = dp[j][k - 1] + abs(n - k - A[i].second) * A[i].first;
        continue;
      }
      if (k == 0) {
        nex[j][k] = dp[j - 1][k] + abs(A[i].second - j + 1) * A[i].first;
        continue;
      }
      nex[j][k] = max(dp[j - 1][k] + abs(A[i].second - j + 1) * A[i].first,
                      dp[j][k - 1] + abs(n - k - A[i].second) * A[i].first);
    }
    for (ll j = 0; j <= i + 1; j++) {
      ll k = i - j + 1;
      dp[j][k] = nex[j][k];
    }
    /* cout << i << endl;
     for(int ii = 0; ii <= i+1; ii++) {
       int j = i - ii + 1;
    cout <<"dp["<<ii<<"]["<<j<<"] = "<< dp[ii][j] << endl;
    }*/
  }

  for (int i = 0; i <= n; i++) {
    int j = n - i;
    ans = max(ans, dp[i][j]);
    // cout <<"dp["<<i<<"]["<<j<<"] = "<< dp[i][j] << endl;
  }
  cout << ans << endl;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <math.h>
#include <vector>

using namespace std;

long long dp[2005][2005] = {};

int main() {

  long long n;
  cin >> n;
  vector<pair<long long, long long>> a(n + 1);
  dp[2005][2005] = {};
  for (long long i = 1; i < n + 1; ++i)
    cin >> a[i].first, a[i].second = i;
  sort(a.begin() + 1, a.end(), greater<pair<long long, long long>>());
  for (long long i = 1; i < n + 1; ++i) {
    for (long long j = 0; j < i; ++j) {
      long long L = j;
      long long R = i - 1 - j;
      dp[L + 1][R] =
          max(dp[L + 1][R], dp[L][R] + abs(L + 1 - a[i].second) * a[i].first);
      R = j;
      L = i - 1 - j;
      dp[L][R + 1] =
          max(dp[L][R + 1], dp[L][R] + abs(n - R - a[i].second) * a[i].first);
    }
  }
  long long ans = 0;
  for (long long i = 1; i < n + 1; ++i) {
    ans = max(ans, dp[i][n - i]);
  }
  cout << ans << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <math.h>
#include <vector>

using namespace std;

long long dp[2005][2005] = {};

int main() {

  long long n;
  cin >> n;
  vector<pair<long long, long long>> a(n + 1);
  for (long long i = 1; i < n + 1; ++i)
    cin >> a[i].first, a[i].second = i;
  sort(a.begin() + 1, a.end(), greater<pair<long long, long long>>());
  for (long long i = 1; i < n + 1; ++i) {
    for (long long j = 0; j < i; ++j) {
      long long L = j;
      long long R = i - 1 - j;
      dp[L + 1][R] =
          max(dp[L + 1][R], dp[L][R] + abs(L + 1 - a[i].second) * a[i].first);
      R = j;
      L = i - 1 - j;
      dp[L][R + 1] =
          max(dp[L][R + 1], dp[L][R] + abs(n - R - a[i].second) * a[i].first);
    }
  }
  long long ans = 0;
  for (long long i = 1; i < n + 1; ++i) {
    ans = max(ans, dp[i][n - i]);
  }
  cout << ans << endl;

  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#define mes(a, b) memset(a, b, sizeof a)
#define pb push_back
using namespace std;
typedef long long ll;
const ll mod = 1e9 + 7;
const int maxn = 2e6 + 10;
const double eps = 1e-1;
ll n;
ll l, r;
ll max(ll x, ll y) { return x > y ? x : y; }
struct Node {
  ll x, id;
  bool operator<(const Node &p) const { return x > p.x; }
} a[maxn];
ll dp[210][2100];
int main() {
  scanf("%lld", &n);
  for (int i = 1; i <= n; i++) {
    scanf("%lld", &a[i].x);
    a[i].id = i;
  }
  sort(a + 1, a + n + 1);
  ll tmp = 0, ans = 0;
  Node b;
  for (int i = 1; i <= n; i++) {
    b = a[i];
    for (ll l = 1, r = n - i + 1; r <= n; l++, r++) {
      dp[i][l] = max(dp[i][l], dp[i - 1][l - 1] + abs(b.id - l) * b.x);
      dp[i][l - 1] = max(dp[i][l - 1], dp[i - 1][l - 1] + abs(b.id - r) * b.x);
    }
  }
  for (int i = 1; i <= n; i++)
    ans = max(ans, dp[n][i]);
  printf("%lld\n", ans);
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#define mes(a, b) memset(a, b, sizeof a)
#define pb push_back
using namespace std;
typedef long long ll;
const ll mod = 1e9 + 7;
const int maxn = 2e6 + 10;
const double eps = 1e-1;
ll n;
ll l, r;
ll max(ll x, ll y) { return x > y ? x : y; }
struct Node {
  ll x, id;
  bool operator<(const Node &p) const { return x > p.x; }
} a[maxn];
ll dp[2100][2100];
int main() {
  scanf("%lld", &n);
  for (int i = 1; i <= n; i++) {
    scanf("%lld", &a[i].x);
    a[i].id = i;
  }
  sort(a + 1, a + n + 1);
  ll tmp = 0, ans = 0;
  Node b;
  for (int i = 1; i <= n; i++) {
    b = a[i];
    for (ll l = 1, r = n - i + 1; r <= n; l++, r++) {
      dp[i][l] = max(dp[i][l], dp[i - 1][l - 1] + abs(b.id - l) * b.x);
      dp[i][l - 1] = max(dp[i][l - 1], dp[i - 1][l - 1] + abs(b.id - r) * b.x);
    }
  }
  for (int i = 1; i <= n; i++)
    ans = max(ans, dp[n][i]);
  printf("%lld\n", ans);
  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL_DEBUG
#include "LOCAL_DEBUG.hpp"
#endif
#define int long long

template <class T> vector<T> make_vec(size_t a) { return vector<T>(a); }
template <class T, class... Ts> auto make_vec(size_t a, Ts... ts) {
  return vector<decltype(make_vec<T>(ts...))>(a, make_vec<T>(ts...));
}
template <class T, class V>
typename enable_if<is_class<T>::value == 0>::type fill(T &t, const V &v) {
  t = v;
}
template <class T, class V>
typename enable_if<is_class<T>::value != 0>::type fill(T &t, const V &v) {
  for (auto &e : t)
    fill(e, v);
}
// auto v = make_vec<int>(h, w);
// fill(v, 0);

using P = pair<int, int>;

signed main() {

  int n;
  cin >> n;
  vector<int> a(n);
  vector<P> p(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    p[i] = {a[i], i};
  }
  sort(p.rbegin(), p.rend());

  auto dp = make_vec<int>(n + 1, n + 1);
  auto dfs = [&](auto &&dfs, int l, int r) -> void {
    if (l + r == n)
      return;
    int i = l + r;
    dp[l + 1][r] = max(dp[l + 1][r], dp[l][r] + p[i].first * (p[i].second - l));
    dp[l][r + 1] =
        max(dp[l][r + 1], dp[l][r] + p[i].first * ((n - 1 - r) - p[i].second));
    dfs(dfs, l + 1, r);
    dfs(dfs, l, r + 1);
  };
  dfs(dfs, 0, 0);

  int ans = 0;
  for (int i = 0; i <= n; i++) {
    for (int j = 0; j <= n; j++) {
      if (i + j == n)
        ans = max(ans, dp[i][j]);
    }
  }
  cout << ans << endl;

  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL_DEBUG
#include "LOCAL_DEBUG.hpp"
#endif
#define int long long

template <class T> vector<T> make_vec(size_t a) { return vector<T>(a); }
template <class T, class... Ts> auto make_vec(size_t a, Ts... ts) {
  return vector<decltype(make_vec<T>(ts...))>(a, make_vec<T>(ts...));
}
template <class T, class V>
typename enable_if<is_class<T>::value == 0>::type fill(T &t, const V &v) {
  t = v;
}
template <class T, class V>
typename enable_if<is_class<T>::value != 0>::type fill(T &t, const V &v) {
  for (auto &e : t)
    fill(e, v);
}
// auto v = make_vec<int>(h, w);
// fill(v, 0);

using P = pair<int, int>;

signed main() {

  int n;
  cin >> n;
  vector<int> a(n);
  vector<P> p(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    p[i] = {a[i], i};
  }
  sort(p.rbegin(), p.rend());

  auto dp = make_vec<int>(n + 1, n + 1);
  for (int i = 0; i < n; i++) {
    for (int j = 0; i + j < n; j++) {
      dp[i + 1][j] =
          max(dp[i + 1][j], dp[i][j] + p[i + j].first * (p[i + j].second - i));
      dp[i][j + 1] =
          max(dp[i][j + 1],
              dp[i][j] + p[i + j].first * ((n - 1 - j) - p[i + j].second));
    }
  }

  int ans = 0;
  for (int i = 0; i <= n; i++) {
    for (int j = 0; j <= n; j++) {
      if (i + j == n)
        ans = max(ans, dp[i][j]);
    }
  }
  cout << ans << endl;

  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef pair<long long, long long> P;
int main() {
  int N;
  cin >> N;
  vector<P> a(N);
  for (int i = 0; i < N; i++) {
    long long b;
    cin >> b;
    a[i] = P(b, i);
  }
  sort(a.begin(), a.end());
  long long dp[N][N];
  long long ans = 0;
  for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j++)
      dp[i][j] = 0;
  }
  dp[0][0] = 0;
  for (int i = 1; i <= N; i++) {
    for (int j = 0; j <= i; j++) {
      int k = i - j;
      if (j == 0)
        dp[j][k] = dp[j][k - 1] + abs(N - k - a[N - i].second) * a[N - i].first;
      else if (k == 0)
        dp[j][k] = dp[j - 1][k] + abs(j - 1 - a[N - i].second) * a[N - i].first;
      else {
        long long m1 =
            dp[j][k - 1] + abs(N - k - a[N - i].second) * a[N - i].first;
        long long m2 =
            dp[j - 1][k] + abs(j - 1 - a[N - i].second) * a[N - i].first;
        dp[j][k] = max(m1, m2);
      }
    }
  }
  for (int i = 0; i <= N; i++)
    ans = max(ans, dp[i][N - i]);
  cout << ans << endl;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef pair<long long, long long> P;
int main() {
  int N;
  cin >> N;
  vector<P> a(N);
  for (int i = 0; i < N; i++) {
    long long b;
    cin >> b;
    a[i] = P(b, i);
  }
  sort(a.begin(), a.end());
  long long dp[N + 1][N + 1];
  long long ans = 0;
  for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j++)
      dp[i][j] = 0;
  }
  dp[0][0] = 0;
  for (int i = 1; i <= N; i++) {
    for (int j = 0; j <= i; j++) {
      int k = i - j;
      if (j == 0)
        dp[j][k] = dp[j][k - 1] + abs(N - k - a[N - i].second) * a[N - i].first;
      else if (k == 0)
        dp[j][k] = dp[j - 1][k] + abs(j - 1 - a[N - i].second) * a[N - i].first;
      else {
        long long m1 =
            dp[j][k - 1] + abs(N - k - a[N - i].second) * a[N - i].first;
        long long m2 =
            dp[j - 1][k] + abs(j - 1 - a[N - i].second) * a[N - i].first;
        dp[j][k] = max(m1, m2);
      }
    }
  }
  for (int i = 0; i <= N; i++)
    ans = max(ans, dp[i][N - i]);
  cout << ans << endl;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <array>
#include <cfenv>
#include <climits>
#include <cmath>
#include <cstdlib>
#include <deque>
#include <forward_list>
#include <fstream>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <ostream>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <streambuf>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define MOD 1000000009
#define SIZE 101
typedef long long ll;
using namespace std;

int main() {
  ll N;
  vector<pair<ll, ll>> A;
  cin >> N;
  for (ll i = 0; i < N; i++) {
    ll tmp;
    cin >> tmp;
    A.push_back({tmp, i});
  }
  sort(A.begin(), A.end());
  reverse(A.begin(), A.end());
  ll DP[SIZE][SIZE];
  for (int i = 0; i < SIZE; i++) {
    for (int j = 0; j < SIZE; j++) {
      DP[i][j] = 0;
    }
  }
  for (int i = 0; i < N; i++) {
    ll act = A[i].first;
    ll no = A[i].second;
    for (int x = 0; x <= i; x++) {
      int y = i - x;
      if (DP[x + 1][y] < DP[x][y] + abs(no - x) * act) {
        DP[x + 1][y] = DP[x][y] + abs(no - x) * act;
      }
      if (DP[x][y + 1] < DP[x][y] + abs(no - (N - 1 - y)) * act) {
        DP[x][y + 1] = DP[x][y] + abs(no - (N - 1 - y)) * act;
      }
    }
  }
  ll ans = 0;
  for (ll i = 0; i <= N; i++) {
    if (ans < DP[i][N - i]) {
      ans = DP[i][N - i];
    }
  }
  cout << ans;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <array>
#include <cfenv>
#include <climits>
#include <cmath>
#include <cstdlib>
#include <deque>
#include <forward_list>
#include <fstream>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <ostream>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <streambuf>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define MOD 1000000009
#define SIZE 2002
typedef long long ll;
using namespace std;

int main() {
  ll N;
  vector<pair<ll, ll>> A;
  cin >> N;
  for (ll i = 0; i < N; i++) {
    ll tmp;
    cin >> tmp;
    A.push_back({tmp, i});
  }
  sort(A.begin(), A.end());
  reverse(A.begin(), A.end());
  ll DP[SIZE][SIZE];
  for (int i = 0; i < SIZE; i++) {
    for (int j = 0; j < SIZE; j++) {
      DP[i][j] = 0;
    }
  }
  for (int i = 0; i < N; i++) {
    ll act = A[i].first;
    ll no = A[i].second;
    for (int x = 0; x <= i; x++) {
      int y = i - x;
      if (DP[x + 1][y] < DP[x][y] + abs(no - x) * act) {
        DP[x + 1][y] = DP[x][y] + abs(no - x) * act;
      }
      if (DP[x][y + 1] < DP[x][y] + abs(no - (N - 1 - y)) * act) {
        DP[x][y + 1] = DP[x][y] + abs(no - (N - 1 - y)) * act;
      }
    }
  }
  ll ans = 0;
  for (ll i = 0; i <= N; i++) {
    if (ans < DP[i][N - i]) {
      ans = DP[i][N - i];
    }
  }
  cout << ans;
  return 0;
}
Problem p02709 (Python)

Buggy submission (status: Time Limit Exceeded):
from heapq import heappop, heappush
from collections import defaultdict

n = int(input())
a = list(map(int, input().split()))

hp = []
for i, e in enumerate(a, 1):
    heappush(hp, [-e, i])

dp = defaultdict(int)
dp[(0, n + 1)] = 0
while hp:
    e, i = heappop(hp)
    e = -e
    dp2 = defaultdict(int)
    for k, v in dp.items():
        l, r = k
        dp2[(l + 1, r)] = max(dp2[(l + 1, r)], v + abs(l + 1 - i) * e)
        dp2[(l, r - 1)] = max(dp2[(l, r - 1)], v + abs(r - 1 - i) * e)

    dp = dp2

ans = 0
for k, v in dp.items():
    ans = max(ans, v)

print(ans)


Fixed submission:
from heapq import heappop, heappush
from collections import defaultdict

n = int(input())
a = list(map(int, input().split()))

hp = []
for i, e in enumerate(a, 1):
    heappush(hp, [-e, i])

dp = defaultdict(int)
dp[(0, n + 1)] = 0
while hp:
    e, i = heappop(hp)
    e = -e
    dp2 = defaultdict(int)
    for k, v in dp.items():
        l, r = k
        dp2[(l + 1, r)] = max(dp2[(l + 1, r)], v + abs(l + 1 - i) * e)
        dp2[(l, r - 1)] = max(dp2[(l, r - 1)], v + abs(r - 1 - i) * e)

    dp = dp2

ans = max(dp.values())

print(ans)
Problem p02709 (C++)

Buggy submission (status: Runtime Error):

// E - Active Infants

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
// const int INF = 2147483647;
// const ll INF = 9223372036854775807;
// const ll MOD = 1e9 + 7;

pair<ll, int> A[2000]; // 活発度とindexのペア
ll dp[2001][2001];     // dp[l][r]:
                   // 左端にl人、右端にr人移動させたときのうれしさの合計の最大値

int main() {
  int N;
  cin >> N;

  for (int i = 0; i < N; i++) {
    ll a;
    cin >> a;
    A[i] = make_pair(a, i);
  }

  sort(A, A + N, greater<pair<ll, int>>());

  ll ans = 0;

  dp[0][0] = 0;

  for (int l = 0; l < N; l++) {
    for (int r = 0; r < N; r++) {
      int i = l + r;
      ll act = A[i].first;
      int loc = A[i].second;

      dp[l + 1][r] = max(dp[l + 1][r], dp[l][r] + act * abs(l - loc));
      dp[l][r + 1] = max(dp[l][r + 1], dp[l][r] + act * abs((N - 1) - r - loc));

      ans = max(ans, dp[l + 1][r]);
      ans = max(ans, dp[l][r + 1]);
    }
  }

  cout << ans << endl;

  return 0;
}

Fixed submission:

// E - Active Infants

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
// const int INF = 2147483647;
// const ll INF = 9223372036854775807;
// const ll MOD = 1e9 + 7;

pair<ll, int> A[2000]; // 活発度とindexのペア
ll dp[2001][2001];     // dp[l][r]:
                   // 左端にl人、右端にr人移動させたときのうれしさの合計の最大値

int main() {
  int N;
  cin >> N;

  for (int i = 0; i < N; i++) {
    ll a;
    cin >> a;
    A[i] = make_pair(a, i);
  }

  sort(A, A + N, greater<pair<ll, int>>());

  ll ans = 0;

  dp[0][0] = 0;

  for (int l = 0; l < N; l++) {
    for (int r = 0; r < N; r++) {
      int i = l + r;
      if (i >= N)
        continue;

      ll act = A[i].first;
      int loc = A[i].second;

      dp[l + 1][r] = max(dp[l + 1][r], dp[l][r] + act * abs(l - loc));
      dp[l][r + 1] = max(dp[l][r + 1], dp[l][r] + act * abs((N - 1) - r - loc));

      ans = max(ans, dp[l + 1][r]);
      ans = max(ans, dp[l][r + 1]);
    }
  }

  cout << ans << endl;

  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
#define rep(i, n) for (ll i = 0; i < (n); ++i)
#define vcout(v)                                                               \
  rep(i, v.size()) cout << v[i] << " ";                                        \
  cout << endl
#define mcout(m)                                                               \
  rep(j, m.size()) { vcout(m[j]); }
const ll MOD = 1e9 + 7;

int mainsolve() {

  ll n;
  cin >> n;
  vector<ll> a(n);
  rep(i, n) cin >> a[i];

  vector<P> ps(n);
  rep(i, n) ps[i] = P(a[i], i);
  sort(ps.rbegin(), ps.rend());

  ll INF = 1e18;
  vector<vector<ll>> dp(n + 1, vector<ll>(n + 1, -INF));
  rep(i, n + 1) dp[0][i] = 0;

  rep(i, n) {
    rep(j, n) {
      ll chooseL = 0, chooseR = 0;
      if (i >= j + 1)
        chooseR = dp[i][j + 1] + ps[i].first * (n - (i - j) - ps[i].second);
      if (i >= j)
        chooseL = dp[i][j] + ps[i].first * (ps[i].second - j);
      if (i + 1 >= j + 1)
        dp[i + 1][j + 1] = max(chooseL, chooseR);
      assert(dp[i + 1][j + 1] >= 0);
    }
  }

  ll res = 0;
  rep(i, n + 1) res = max(res, dp[n][i]);
  cout << res << endl;
  return 0;
}

int main() { return mainsolve(); }

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
#define rep(i, n) for (ll i = 0; i < (n); ++i)
#define vcout(v)                                                               \
  rep(i, v.size()) cout << v[i] << " ";                                        \
  cout << endl
#define mcout(m)                                                               \
  rep(j, m.size()) { vcout(m[j]); }
const ll MOD = 1e9 + 7;

int mainsolve() {

  ll n;
  cin >> n;
  vector<ll> a(n);
  rep(i, n) cin >> a[i];

  vector<P> ps(n);
  rep(i, n) ps[i] = P(a[i], i);
  sort(ps.rbegin(), ps.rend());

  ll INF = 1e18;
  vector<vector<ll>> dp(n + 1, vector<ll>(n + 1, -INF));
  rep(i, n + 1) dp[0][i] = 0;

  rep(i, n) {
    rep(j, n + 1) {
      if (i < j)
        continue;
      dp[i + 1][j] =
          max(dp[i + 1][j],
              dp[i][j] + ps[i].first * (n - (i - j) - 1 - ps[i].second));
      if (j + 1 <= n)
        dp[i + 1][j + 1] =
            max(dp[i + 1][j + 1], dp[i][j] + ps[i].first * (ps[i].second - j));
    }
  }

  ll res = 0;
  rep(i, n + 1) res = max(res, dp[n][i]);
  cout << res << endl;
  return 0;
}

int main() { return mainsolve(); }
Problem p02709 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <iostream>
#include <stdio.h>
#include <vector>
#define all(x) x.begin(), x.end()
using namespace std;
typedef long long ll;
typedef double db;
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
  int n;
  cin >> n;
  vector<pair<int, int>> a(n);
  for (int i = 0; i < n; ++i) {
    cin >> a[i].first;
    a[i].second = i;
  }
  sort(all(a));
  reverse(all(a));
  vector<vector<ll>> dp(n + 1, vector<ll>(n + 1, 0));
  dp[0][0] = 0;
  for (int i = 0; i < n; ++i) {
    for (int j = 1; j <= i + 1; ++j) {
      dp[j][i + 1 - j] = max(dp[j][i + 1 - j],
                             dp[j - 1][i + 1 - j] +
                                 1ll * a[i].first * abs(a[i].second + 1 - j));
      dp[i + 1 - j][j] = max(dp[i + 1 - j][j],
                             dp[i + 1 - j][j - 1] +
                                 1ll * a[i].first * abs(a[i].second - (n - j)));
    }
  }
  ll ans = 0;
  for (int i = 0; i <= n; ++i)
    ans = max(ans, dp[i][n - i]);
  cout << ans;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <stdio.h>
#include <vector>
#define all(x) x.begin(), x.end()
using namespace std;
typedef long long ll;
typedef double db;
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  // freopen("input.txt", "r", stdin);
  // freopen("output.txt", "w", stdout);
  int n;
  cin >> n;
  vector<pair<int, int>> a(n);
  for (int i = 0; i < n; ++i) {
    cin >> a[i].first;
    a[i].second = i;
  }
  sort(all(a));
  reverse(all(a));
  vector<vector<ll>> dp(n + 1, vector<ll>(n + 1, 0));
  dp[0][0] = 0;
  for (int i = 0; i < n; ++i) {
    for (int j = 1; j <= i + 1; ++j) {
      dp[j][i + 1 - j] = max(dp[j][i + 1 - j],
                             dp[j - 1][i + 1 - j] +
                                 1ll * a[i].first * abs(a[i].second + 1 - j));
      dp[i + 1 - j][j] = max(dp[i + 1 - j][j],
                             dp[i + 1 - j][j - 1] +
                                 1ll * a[i].first * abs(a[i].second - (n - j)));
    }
  }
  ll ans = 0;
  for (int i = 0; i <= n; ++i)
    ans = max(ans, dp[i][n - i]);
  cout << ans;
  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <stdio.h>
using namespace std;

const double pi = acos(-1);

typedef long long ll;
const int maxn = 2005;
struct pp {
  ll v;
  int id;
  bool operator<(const pp &o) const {
    if (v != o.v)
      return v > o.v;
    return id < o.id;
  }
};
pp A[maxn];
int mark[maxn];
ll dp[maxn][maxn];

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(0);
#ifdef acm
  freopen("a.in", "r", stdin);
#endif // acm

  int n;
  cin >> n;

  for (int i = 1; i <= n; i++)
    cin >> A[i].v, A[i].id = i;

  sort(A + 1, A + n + 1);

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

  dp[0][0] = 0;

  ll sol = 0;
  for (int i = 0; i <= n; i++) {
    for (int j = 0; j <= n; j++) {
      if (dp[i][j] == -1)
        continue;
      int t = i + j + 1;
      int ini = i + 1;
      int fin = n - j;
      dp[i + 1][j] = max(dp[i + 1][j], dp[i][j] + abs(A[t].id - ini) * A[t].v);
      dp[i][j + 1] = max(dp[i][j + 1], dp[i][j] + abs(A[t].id - fin) * A[t].v);
    }
  }

  for (int i = 0; i <= n; i++)
    sol = max(sol, dp[i][n - i]);

  cout << sol << "\n";

  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <stdio.h>
using namespace std;

const double pi = acos(-1);

typedef long long ll;
const int maxn = 2005;
struct pp {
  ll v;
  int id;
  bool operator<(const pp &o) const {
    if (v != o.v)
      return v > o.v;
    return id < o.id;
  }
};
pp A[maxn];
int mark[maxn];
ll dp[maxn][maxn];

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(0);
#ifdef acm
  freopen("a.in", "r", stdin);
#endif // acm

  int n;
  cin >> n;

  for (int i = 1; i <= n; i++)
    cin >> A[i].v, A[i].id = i;

  sort(A + 1, A + n + 1);

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

  dp[0][0] = 0;

  ll sol = 0;
  for (int i = 0; i <= n; i++) {
    for (int j = 0; j <= n; j++) {
      if (i + j > n)
        continue;
      if (dp[i][j] == -1)
        continue;
      int t = i + j + 1;
      int ini = i + 1;
      int fin = n - j;
      dp[i + 1][j] = max(dp[i + 1][j], dp[i][j] + abs(A[t].id - ini) * A[t].v);
      dp[i][j + 1] = max(dp[i][j + 1], dp[i][j] + abs(A[t].id - fin) * A[t].v);
    }
  }

  for (int i = 0; i <= n; i++)
    sol = max(sol, dp[i][n - i]);

  cout << sol << "\n";

  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#define simplify                                                               \
  ios_base::sync_with_stdio(0);                                                \
  cin.tie(0);                                                                  \
  cout.tie(0);
#define int long long
#define endl "\n"
using namespace std;

vector<pair<int, int>> arr;
int dp[2005][2005];
int n;
int solve(int pos, int left, int right) {
  if (pos >= n)
    return 0;
  if (left > right)
    return -1e10;
  int &ans = dp[left][right];

  if (ans != -1)
    return ans;
  int maxi1 = solve(pos + 1, left + 1, right) +
              abs(arr[pos].second - left) * arr[pos].first;
  int maxi2 = solve(pos + 1, left, right - 1) +
              abs(arr[pos].second - right) * arr[pos].first;

  return max(maxi1, maxi2);
}

int32_t main() {

  simplify cin >> n;
  arr.resize(n);
  for (int i = 0; i < n; i++) {
    cin >> arr[i].first;
    arr[i].second = i;
  }
  sort(arr.rbegin(), arr.rend());
  memset(dp, -1, sizeof dp);
  int ans = solve(0, 0, n - 1);
  cout << ans;

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#define simplify                                                               \
  ios_base::sync_with_stdio(0);                                                \
  cin.tie(0);                                                                  \
  cout.tie(0);
#define int long long
#define endl "\n"
using namespace std;

vector<pair<int, int>> arr;
int dp[2005][2005];
int n;
int solve(int pos, int left, int right) {
  if (pos >= n)
    return 0;
  if (left > right)
    return -1e10;
  int &ans = dp[left][right];

  if (ans != -1)
    return ans;
  int maxi1 = solve(pos + 1, left + 1, right) +
              abs(arr[pos].second - left) * arr[pos].first;
  int maxi2 = solve(pos + 1, left, right - 1) +
              abs(arr[pos].second - right) * arr[pos].first;

  return ans = max(maxi1, maxi2);
}

int32_t main() {

  simplify cin >> n;
  arr.resize(n);
  for (int i = 0; i < n; i++) {
    cin >> arr[i].first;
    arr[i].second = i;
  }
  sort(arr.rbegin(), arr.rend());
  memset(dp, -1, sizeof dp);
  int ans = solve(0, 0, n - 1);
  cout << ans;

  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

#define int long long
#define F first
#define S second
#define fast ios::sync_with_stdio(false), cin.tie(nullptr);
const int N = 1005;
int dp[N][N];
int n;
vector<pair<int, int>> a;
int solve(int l, int r) {
  if (l > r)
    return 0;
  int &res = dp[l][r];
  if (res != -1)
    return res;
  int id = n - (r - l + 1);
  int pos = a[id].S;
  int val = a[id].F;
  res = max(res, solve(l + 1, r) + abs(pos - l) * val);
  res = max(res, solve(l, r - 1) + abs(pos - r) * val);
  return res;
}
int32_t main() {
  fast;
  cin >> n;
  a.resize(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i].F;
    a[i].S = i;
  }
  sort(a.rbegin(), a.rend());
  memset(dp, -1, sizeof(dp));
  cout << solve(0, n - 1);
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;

#define int long long
#define F first
#define S second
#define fast ios::sync_with_stdio(false), cin.tie(nullptr);
const int N = 2005;
int dp[N][N];
int n;
vector<pair<int, int>> a;
int solve(int l, int r) {
  if (l > r)
    return 0;
  int &res = dp[l][r];
  if (res != -1)
    return res;
  int id = n - (r - l + 1);
  int pos = a[id].S;
  int val = a[id].F;
  res = max(res, solve(l + 1, r) + abs(pos - l) * val);
  res = max(res, solve(l, r - 1) + abs(pos - r) * val);
  return res;
}
int32_t main() {
  fast;
  cin >> n;
  a.resize(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i].F;
    a[i].S = i;
  }
  sort(a.rbegin(), a.rend());
  memset(dp, -1, sizeof(dp));
  cout << solve(0, n - 1);
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define fs first
#define sc second
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define ALL(A) A.begin(), A.end()
#define RALL(A) A.rbegin(), A.rend()
typedef long long LL;
typedef pair<LL, LL> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> T gcd(T a, T b) { return b ? gcd(b, a % b) : a; }
const LL mod = 1000000007;
const LL LINF = 1LL << 62;
const int INF = 1 << 30;
int dx[] = {1, 0, -1, 0, 1, -1, 1, -1};
int dy[] = {0, 1, 0, -1, 1, -1, -1, 1};

int main() {
  int n;
  cin >> n;
  vector<P> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i].fs;
    a[i].sc = i;
  }
  sort(RALL(a));
  vector<vector<LL>> dp(n + 1, vector<LL>(n + 1, -LINF));
  dp[0][n - 1] = 0;
  for (int i = 0; i < n; i++) {
    for (int l = 0; l <= i; l++) {
      int r = n - (i - l) - 1;
      chmax(dp[l + 1][r], dp[l][r] + a[i].fs * abs(a[i].sc - l));
      chmax(dp[l][r - 1], dp[l][r] + a[i].fs * abs(a[i].sc - r));
    }
  }
  LL ans = 0;
  for (int i = 1; i <= n; i++) {
    chmax(ans, dp[i][i - 1]);
    chmax(ans, dp[i - 1][i]);
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define fs first
#define sc second
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define ALL(A) A.begin(), A.end()
#define RALL(A) A.rbegin(), A.rend()
typedef long long LL;
typedef pair<LL, LL> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> T gcd(T a, T b) { return b ? gcd(b, a % b) : a; }
const LL mod = 1000000007;
const LL LINF = 1LL << 62;
const int INF = 1 << 30;
int dx[] = {1, 0, -1, 0, 1, -1, 1, -1};
int dy[] = {0, 1, 0, -1, 1, -1, -1, 1};

int main() {
  int n;
  cin >> n;
  vector<P> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i].fs;
    a[i].sc = i;
  }
  sort(RALL(a));
  vector<vector<LL>> dp(n + 1, vector<LL>(n + 1, -LINF));
  dp[0][n - 1] = 0;
  for (int i = 0; i < n; i++) {
    for (int l = 0; l <= i; l++) {
      int r = n - (i - l) - 1;
      if (l + 1 <= n)
        chmax(dp[l + 1][r], dp[l][r] + a[i].fs * abs(a[i].sc - l));
      if (r - 1 >= 0)
        chmax(dp[l][r - 1], dp[l][r] + a[i].fs * abs(a[i].sc - r));
    }
  }
  LL ans = 0;
  for (int i = 1; i <= n; i++) {
    chmax(ans, dp[i][i - 1]);
    chmax(ans, dp[i - 1][i]);
  }
  cout << ans << endl;
  return 0;
}
Problem p02709 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; i++)
#define rrep(i, n) for (int i = n - 1; i >= 0; i--)
using namespace std;
#define INF ((1 << 30) - 1)
#define LINF (1LL << 60)
#define EPS (1e-10)
typedef long long ll;
typedef pair<ll, ll> P;
const int MOD = 1000000007;
const int MOD2 = 998244353;

vector<P> vp;
ll dp[1010][1010];

int main() {
  ll n;
  cin >> n;
  rep(i, n) {
    ll a;
    cin >> a;
    vp.push_back({a, i});
  }
  sort(vp.begin(), vp.end());
  reverse(vp.begin(), vp.end());

  rep(i, n + 1) {
    for (int j = 0; j + i <= n; j++) {
      ll cl = vp[i + j].first * (vp[i + j].second - i);
      ll cr = vp[i + j].first * (n - j - 1 - vp[i + j].second);
      dp[i + 1][j] = max(dp[i + 1][j], dp[i][j] + cl);
      dp[i][j + 1] = max(dp[i][j + 1], dp[i][j] + cr);
    }
  }

  ll ans = 0;
  rep(i, n) ans = max(ans, dp[i][n - i]);

  cout << ans << endl;

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; i++)
#define rrep(i, n) for (int i = n - 1; i >= 0; i--)
using namespace std;
#define INF ((1 << 30) - 1)
#define LINF (1LL << 60)
#define EPS (1e-10)
typedef long long ll;
typedef pair<ll, ll> P;
const int MOD = 1000000007;
const int MOD2 = 998244353;

vector<P> vp;
ll dp[2010][2010];

int main() {
  ll n;
  cin >> n;
  rep(i, n) {
    ll a;
    cin >> a;
    vp.push_back({a, i});
  }
  sort(vp.begin(), vp.end());
  reverse(vp.begin(), vp.end());

  rep(i, n + 1) {
    for (int j = 0; j + i <= n; j++) {
      ll cl = vp[i + j].first * (vp[i + j].second - i);
      ll cr = vp[i + j].first * (n - j - 1 - vp[i + j].second);
      dp[i + 1][j] = max(dp[i + 1][j], dp[i][j] + cl);
      dp[i][j + 1] = max(dp[i][j + 1], dp[i][j] + cr);
    }
  }

  ll ans = 0;
  rep(i, n) ans = max(ans, dp[i][n - i]);

  cout << ans << endl;

  return 0;
}
Problem p02710 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

template <typename T> void out(T x) {
  cout << x << endl;
  exit(0);
}
#define watch(x) cout << (#x) << " is " << (x) << endl

using ll = long long;

const int maxn = 1e6 + 5;

int n;
int c[maxn];
vector<int> g[maxn];
int siz[maxn];

ll ans[maxn];

void dfs(int at, int p, map<ll, ll> &mp) {
  siz[at] = 1;
  int C = c[at];
  for (int to : g[at]) {
    if (to == p)
      continue;
    map<ll, ll> mp_to;
    dfs(to, at, mp_to);
    siz[at] += siz[to];
    ll middle = siz[to] - mp_to[C];
    ans[C] -= (middle * (middle + 1) / 2);
    for (auto pr : mp_to) {
      mp[pr.first] += pr.second;
    }
  }

  mp[C] = siz[at];
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  for (int i = 1; i <= n; i++) {
    ans[i] = 1ll * n * (n + 1) / 2;
  }

  map<ll, ll> mp;
  dfs(1, 0, mp);

  for (int i = 1; i <= n; i++) {
    ll top = n - mp[i];
    ans[i] -= (top * (top + 1) / 2);
  }

  for (int i = 1; i <= n; i++) {
    cout << ans[i] << "\n";
  }

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

template <typename T> void out(T x) {
  cout << x << endl;
  exit(0);
}
#define watch(x) cout << (#x) << " is " << (x) << endl

using ll = long long;

const int maxn = 1e6 + 5;

int n;
int c[maxn];
vector<int> g[maxn];
int siz[maxn];

ll ans[maxn];

void dfs(int at, int p, map<ll, ll> &mp) {
  siz[at] = 1;
  int C = c[at];
  for (int to : g[at]) {
    if (to == p)
      continue;
    map<ll, ll> mp_to;
    dfs(to, at, mp_to);
    siz[at] += siz[to];
    ll middle = siz[to] - mp_to[C];
    ans[C] -= (middle * (middle + 1) / 2);
    if (mp.size() < mp_to.size()) {
      swap(mp, mp_to);
    }
    for (auto pr : mp_to) {
      mp[pr.first] += pr.second;
    }
  }

  mp[C] = siz[at];
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  for (int i = 1; i <= n; i++) {
    ans[i] = 1ll * n * (n + 1) / 2;
  }

  map<ll, ll> mp;
  dfs(1, 0, mp);

  for (int i = 1; i <= n; i++) {
    ll top = n - mp[i];
    ans[i] -= (top * (top + 1) / 2);
  }

  for (int i = 1; i <= n; i++) {
    cout << ans[i] << "\n";
  }

  return 0;
}
Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#include <iostream>
using namespace std;
const long long MOD = 1e9 + 7;
const long long MAXN = 5e3 + 5;
long long n;
vector<long long> v1[MAXN];
long long ans[MAXN];
long long res[MAXN];
long long col[MAXN];

long long sz[MAXN];
long long f(long long x) { return 1LL * x * (x + 1) / 2; }
void dfs(long long curr, long long par) {
  long long tmp = res[col[curr]];
  for (long long x : v1[curr]) {
    if (x != par) {
      res[col[curr]] = 0;
      dfs(x, curr);
      sz[curr] += sz[x];
      ans[col[curr]] -= f(sz[x] - res[col[curr]]);
    }
  }
  res[col[curr]] = tmp + sz[curr];
}
int main() {
  cin >> n;
  for (long long i = 1; i <= n; i++) {
    cin >> col[i];
    sz[i] = 1;
  }
  for (long long i = 1; i < n; i++) {
    long long x, y;
    cin >> x >> y;
    v1[x].push_back(y);
    v1[y].push_back(x);
  }
  dfs(1, 1);
  for (long long i = 1; i <= n; i++) {
    // cout<<col[i]<<endl;
    cout << (ans[i] + f(n) - f(n - res[i])) << endl;
  }
}

Fixed submission:
#include <bits/stdc++.h>
#include <iostream>
using namespace std;
const long long MOD = 1e9 + 7;
const long long MAXN = 2e5 + 5;
long long n;
vector<long long> v1[MAXN];
long long ans[MAXN];
long long res[MAXN];
long long col[MAXN];

long long sz[MAXN];
long long f(long long x) { return 1LL * x * (x + 1) / 2; }
void dfs(long long curr, long long par) {
  long long tmp = res[col[curr]];
  for (long long x : v1[curr]) {
    if (x != par) {
      res[col[curr]] = 0;
      dfs(x, curr);
      sz[curr] += sz[x];
      ans[col[curr]] -= f(sz[x] - res[col[curr]]);
    }
  }
  res[col[curr]] = tmp + sz[curr];
}
int main() {
  cin >> n;
  for (long long i = 1; i <= n; i++) {
    cin >> col[i];
    sz[i] = 1;
  }
  for (long long i = 1; i < n; i++) {
    long long x, y;
    cin >> x >> y;
    v1[x].push_back(y);
    v1[y].push_back(x);
  }
  dfs(1, 1);
  for (long long i = 1; i <= n; i++) {
    // cout<<col[i]<<endl;
    cout << (ans[i] + f(n) - f(n - res[i])) << endl;
  }
}
Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int MAX = 510000;
const int MOD = 1000000007;
int sz[200001], zen[200001], c[200001], ans[200001];
bool used[200001];
vector<int> hen[200000];
int dfssz(int a) {

  sz[a] = 1;
  for (auto &&i : hen[a])
    if (!sz[i])
      sz[a] += dfssz(i);

  return sz[a];
}
void dfs(int a) {
  used[a] = true;
  ans[c[a]] += 1 + zen[c[a]] * (zen[c[a]] - 1) / 2;
  zen[c[a]] -= sz[a];
  ans[c[a]] -= (zen[c[a]] - 1) * zen[c[a]] / 2;
  int tmp = zen[c[a]];
  for (auto &&i : hen[a])
    if (!used[i]) {

      ans[c[a]] -= sz[i] * (sz[i] - 1) / 2;
      zen[c[a]] = sz[i];
      dfs(i);
    }
  zen[c[a]] = tmp;
  return;
}

signed main() {
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }
  int a, b;
  for (int i = 0; i < n - 1; i++) {
    cin >> a >> b;
    hen[a].push_back(b);
    hen[b].push_back(a);
  }
  fill(ans, ans + n + 1, 0);
  fill(sz, sz + n + 1, 0);
  fill(zen, zen + n + 1, n);
  dfssz(1);
  fill(used, used + n + 1, false);
  dfs(1);
  for (int i = 1; i <= n; i++)
    cout << ans[i] << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int MAX = 510000;
const int MOD = 1000000007;
int sz[200001], zen[200001], c[200001], ans[200001];
bool used[200001];
vector<int> hen[200001];
int dfssz(int a) {

  sz[a] = 1;
  for (auto &&i : hen[a])
    if (!sz[i])
      sz[a] += dfssz(i);

  return sz[a];
}
void dfs(int a) {
  used[a] = true;
  ans[c[a]] += 1 + zen[c[a]] * (zen[c[a]] - 1) / 2;
  zen[c[a]] -= sz[a];
  ans[c[a]] -= (zen[c[a]] - 1) * zen[c[a]] / 2;
  int tmp = zen[c[a]];
  for (auto &&i : hen[a])
    if (!used[i]) {

      ans[c[a]] -= sz[i] * (sz[i] - 1) / 2;
      zen[c[a]] = sz[i];
      dfs(i);
    }
  zen[c[a]] = tmp;
  return;
}

signed main() {
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }
  int a, b;
  for (int i = 0; i < n - 1; i++) {
    cin >> a >> b;
    hen[a].push_back(b);
    hen[b].push_back(a);
  }
  fill(ans, ans + n + 1, 0);
  fill(sz, sz + n + 1, 0);
  fill(zen, zen + n + 1, n);
  dfssz(1);
  fill(used, used + n + 1, false);
  dfs(1);
  for (int i = 1; i <= n; i++)
    cout << ans[i] << endl;
}
Problem p02710 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cctype>
#include <cfloat>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <vector>
using namespace std;
typedef long long LL;
typedef pair<LL, LL> P;

const int L = 2e5;
LL c[L];
vector<int> es[L];
LL ps[L];
LL w[L];
LL cnt[L];
LL ret[L];
LL N;
LL dfs(int r, int p) {
  int pcnt = cnt[c[r]];
  int pps = ps[c[r]];
  ps[c[r]] = r;
  LL s = 1;
  ret[c[r]] += s;
  LL sum = 1;
  for (auto i : es[r]) {
    if (i == p) {
      continue;
    }
    w[r] = 0;
    auto v = dfs(i, r);
    auto t = w[r] + v;
    ret[c[r]] += t * s;
    s += t;
    sum += v;
  }
  {
    auto t = N - sum - pcnt;
    ret[c[r]] += t * s;
  }
  cnt[c[r]] += s;
  w[pps] -= sum;
  ps[c[r]] = pps;
  return sum;
}

int main() {
  scanf("%lld", &N);
  for (int i = 1; i <= N; i++) {
    scanf("%lld", &c[i]);
  }
  for (int i = 0; i < N - 1; i++) {
    LL a, b;
    scanf("%lld", &a);
    scanf("%lld", &b);
    es[a].push_back(b);
    es[b].push_back(a);
  }
  dfs(1, 0);
  for (int i = 1; i <= N; i++) {
    cout << ret[i] << endl;
  }
  return 0;
}


Fixed submission:
#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cctype>
#include <cfloat>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <vector>
using namespace std;
typedef long long LL;
typedef pair<LL, LL> P;

const int L = 2e5 + 9;
LL c[L];
vector<int> es[L];
LL ps[L];
LL w[L];
LL cnt[L];
LL ret[L];
LL N;
LL dfs(int r, int p) {
  int pcnt = cnt[c[r]];
  int pps = ps[c[r]];
  ps[c[r]] = r;
  LL s = 1;
  ret[c[r]] += s;
  LL sum = 1;
  for (auto i : es[r]) {
    if (i == p) {
      continue;
    }
    w[r] = 0;
    auto v = dfs(i, r);
    auto t = w[r] + v;
    ret[c[r]] += t * s;
    s += t;
    sum += v;
  }
  {
    auto t = N - sum - pcnt;
    ret[c[r]] += t * s;
  }
  cnt[c[r]] += s;
  w[pps] -= sum;
  ps[c[r]] = pps;
  return sum;
}

int main() {
  scanf("%lld", &N);
  for (int i = 1; i <= N; i++) {
    scanf("%lld", &c[i]);
  }
  for (int i = 0; i < N - 1; i++) {
    LL a, b;
    scanf("%lld", &a);
    scanf("%lld", &b);
    es[a].push_back(b);
    es[b].push_back(a);
  }
  dfs(1, 0);
  for (int i = 1; i <= N; i++) {
    cout << ret[i] << endl;
  }
  return 0;
}