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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <set>
#include <string>
#include <vector>

using namespace std;

class IndexedActivity {
public:
  int index;
  long long activity;

  IndexedActivity() : index(0), activity(0LL) {}

  IndexedActivity(int index_, long long activity_)
      : index(index_), activity(activity_) {}

  bool operator<(const IndexedActivity &rhs) const {
    return activity < rhs.activity;
  }
  bool operator>(const IndexedActivity &rhg) const {
    return activity > rhg.activity;
  }
};

long long maxFromNew(size_t forwardInfants, size_t backwardInfants,
                     const vector<IndexedActivity> &sortedActivities,
                     vector<vector<long long>> &memo) {
  size_t N = sortedActivities.size();
  if (forwardInfants + backwardInfants == N) {
    return 0LL;
  }
  if (memo[forwardInfants][backwardInfants] >= 0) {
    return memo[forwardInfants][backwardInfants];
  }

  const IndexedActivity &nextInfant =
      sortedActivities[forwardInfants + backwardInfants];

  int forwardMove = nextInfant.index - forwardInfants;
  long long letNextGoForward =
      (forwardMove >= 0) ? maxFromNew(forwardInfants + 1, backwardInfants,
                                      sortedActivities, memo) +
                               nextInfant.activity * forwardMove
                         : 0LL;

  int backwardMove = (N - 1 - backwardInfants) - nextInfant.index;
  long long letNextGoBackward =
      (backwardMove >= 0)
          ? maxFromNew(forwardInfants, backwardInfants + 1, sortedActivities,
                       memo) +
                nextInfant.activity *
                    ((N - 1 - backwardInfants) - nextInfant.index)
          : 0LL;

  long long ret = max(letNextGoForward, letNextGoBackward);
  memo[forwardInfants][backwardInfants] = ret;
  return ret;
}

int main() {
  int N;
  cin >> N;
  vector<long long> A(N, 0LL);
  for (int i = 0; i < N; ++i) {
    cin >> A[i];
  }

  // bitset<2000> path(0);
  // long long ans = maxFrom(0, path, A);
  // long long ans = greedy(A);
  vector<IndexedActivity> activities(A.size(), IndexedActivity());
  vector<vector<long long>> memo(A.size(), vector<long long>(A.size(), -1));
  for (size_t i = 0; i < A.size(); ++i) {
    activities[i] = IndexedActivity(i, A[i]);
  }
  sort(activities.begin(), activities.end(), greater<IndexedActivity>());
  long long ans = maxFromNew(0, 0, activities, memo);
  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <set>
#include <string>
#include <vector>

using namespace std;

long long maxFrom(int currIndex, bitset<2000> path,
                  const vector<long long> &A) {
  int N = (int)A.size();

  if (currIndex == N - 1) {
    for (int i = 0; i < N; ++i) {
      if (!path.test(i)) {
        return A[currIndex] * abs(currIndex - i);
      }
    }
  }

  long long currMax = -1LL;
  bitset<2000> currPath = path;
  for (int i = 0; i < N; ++i) {
    if (path.test(i)) {
      continue;
    }
    bitset<2000> tempPath = path;
    tempPath.set(i);
    long long tempMax =
        maxFrom(currIndex + 1, tempPath, A) + A[currIndex] * abs(currIndex - i);
    if (currMax < 0 || currMax < tempMax) {
      currMax = tempMax;
      currPath = tempPath;
    }
  }
  return currMax;
}

class IndexedActivity {
public:
  int index;
  long long activity;

  IndexedActivity() : index(0), activity(0LL) {}

  IndexedActivity(int index_, long long activity_)
      : index(index_), activity(activity_) {}

  bool operator<(const IndexedActivity &rhs) const {
    return activity < rhs.activity;
  }
  bool operator>(const IndexedActivity &rhg) const {
    return activity > rhg.activity;
  }
};

long long greedy(const vector<long long> &A) {
  vector<IndexedActivity> activities;
  for (int i = 0; i < A.size(); ++i) {
    activities.push_back(IndexedActivity(i, A[i]));
  }
  sort(activities.begin(), activities.end(), greater<IndexedActivity>());
  int left = 0;
  int right = A.size() - 1;
  long long happiness = 0LL;
  for (int i = 0; i < activities.size(); ++i) {
    int moveToLeft = abs(activities[i].index - left);
    int moveToRight = abs(activities[i].index - right);
    if (moveToLeft < moveToRight) {
      happiness += moveToRight * activities[i].activity;
      --right;
    } else {
      happiness += moveToLeft * activities[i].activity;
      ++left;
    }
  }
  return happiness;
}

long long maxFromNew(int forwardInfants, int backwardInfants,
                     const vector<IndexedActivity> &sortedActivities) {
  int N = sortedActivities.size();
  if (forwardInfants + backwardInfants == N) {
    return 0LL;
  }
  const IndexedActivity &nextInfant =
      sortedActivities[forwardInfants + backwardInfants];
  long long letNextGoForward = -1LL;
  int forwardMove = nextInfant.index - forwardInfants;
  if (forwardMove >= 0) {
    letNextGoForward =
        maxFromNew(forwardInfants + 1, backwardInfants, sortedActivities) +
        nextInfant.activity * forwardMove;
  }
  long long letNextGoBackward = -1LL;
  int backwardMove = (N - 1 - backwardInfants) - nextInfant.index;
  if (backwardMove >= 0) {
    letNextGoBackward =
        maxFromNew(forwardInfants, backwardInfants + 1, sortedActivities) +
        nextInfant.activity * ((N - 1 - backwardInfants) - nextInfant.index);
  }
  return max(letNextGoForward, letNextGoBackward);
}

int main() {
  int N;
  cin >> N;
  vector<long long> A(N, 0LL);
  for (int i = 0; i < N; ++i) {
    cin >> A[i];
  }

  // bitset<2000> path(0);
  // long long ans = maxFrom(0, path, A);
  // long long ans = greedy(A);
  vector<IndexedActivity> activities(A.size(), IndexedActivity());
  for (int i = 0; i < A.size(); ++i) {
    activities[i] = IndexedActivity(i, A[i]);
  }
  sort(activities.begin(), activities.end(), greater<IndexedActivity>());
  long long ans = maxFromNew(0, 0, activities);
  cout << ans << endl;
  return 0;
}