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


#include <bits/stdc++.h>

#define int long long int
#define rep(a, b, c) for (int a = b; a < c; a++)
#define repm(a, b, c) for (int a = (b - 1); a >= c; a--)
#define pb push_back
#define str string
#define sf(a) scanfs("%d", &a)
#define pb push_back
#define mp make_pair
#define Fi first
#define Se second
#define ALL(v) (v).begin(), (v).end()

using namespace std;

const int INF = 1e18 + 9;
const int Mod = 1e9 + 7;
inline int replac(str s) {
  double ans = 0;
  rep(i, 0, s.length()) { ans += (s[i] - '0') * pow(10, s.length() - i - 1); }
  return (int)ans;
}
inline string numstr(int m) {
  str s = "";
  while (m > 0) {
    char c;
    int a = m / 10;
    if (a > 0)
      a = m % (a * 10);
    else
      a = m;
    c = (char)('0' + a);
    s += c;
    m /= 10;
  }
  str st = "";
  for (int i = s.size() - 1; i >= 0; i--) {
    st += s[i];
  }
  return st;
}
typedef pair<int, int> pii;
typedef vector<pii> vii;
int n, k, memo[10][2001];
vector<int> book[10];
int f(int nu = 0, int cou = 0) {
  if (nu >= 10 || cou >= k) {
    return 0;
  } else if (memo[nu][cou] > 0) {
    return memo[nu][cou];
  } else {
    int ans = f(nu + 1, cou);
    rep(i, 0, book[nu].size()) {
      if ((i + cou) >= k)
        break;
      ans = max(ans, f(nu + 1, cou + i + 1) + book[nu][i]);
    }
    return memo[nu][cou] = ans;
  }
}
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cin >> n >> k;
  rep(i, 0, 10) {
    rep(j, 0, 2001) { memo[i][j] = 0; }
  }
  rep(i, 0, n) {
    int x, y;
    cin >> x >> y;
    book[y - 1].pb(x);
  }
  rep(i, 0, 10) {
    if (!book[i].empty())
      sort(ALL(book[i]), greater<int>());
  }
  rep(i, 0, 10) {
    rep(j, 1, book[i].size()) {
      book[i][j] += book[i][j - 1];
      // cout << i << " "<<book[i][j] << endl;
    }
  }
  rep(i, 0, 10) {
    rep(j, 0, book[i].size()) { book[i][j] += j * (j + 1); }
  }
  cout << f() << endl;
  // cout << book[1][3] << endl;
  return 0;
}

Error: TLE

Fixed submission:


#include <bits/stdc++.h>

#define int long long int
#define rep(a, b, c) for (int a = b; a < c; a++)
#define repm(a, b, c) for (int a = (b - 1); a >= c; a--)
#define pb push_back
#define str string
#define sf(a) scanfs("%d", &a)
#define pb push_back
#define mp make_pair
#define Fi first
#define Se second
#define ALL(v) (v).begin(), (v).end()

using namespace std;

const int INF = 1e18 + 9;
const int Mod = 1e9 + 7;
inline int replac(str s) {
  double ans = 0;
  rep(i, 0, s.length()) { ans += (s[i] - '0') * pow(10, s.length() - i - 1); }
  return (int)ans;
}
inline string numstr(int m) {
  str s = "";
  while (m > 0) {
    char c;
    int a = m / 10;
    if (a > 0)
      a = m % (a * 10);
    else
      a = m;
    c = (char)('0' + a);
    s += c;
    m /= 10;
  }
  str st = "";
  for (int i = s.size() - 1; i >= 0; i--) {
    st += s[i];
  }
  return st;
}
typedef pair<int, int> pii;
typedef vector<pii> vii;
int n, k;
vector<int> book[10];
int f(int nu = 0, int cou = 0) {
  if (nu >= 10 || cou >= k) {
    return 0;
  } else {
    int ans = f(nu + 1, cou);
    rep(i, 0, book[nu].size()) {
      if ((i + cou) >= k)
        break;
      ans = max(ans, f(nu + 1, cou + i + 1) + book[nu][i]);
    }
    return ans;
  }
}
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cin >> n >> k;
  rep(i, 0, n) {
    int x, y;
    cin >> x >> y;
    book[y - 1].pb(x);
  }
  rep(i, 0, 10) {
    if (!book[i].empty())
      sort(ALL(book[i]), greater<int>());
  }
  rep(i, 0, 10) {
    rep(j, 1, book[i].size()) {
      book[i][j] += book[i][j - 1];
      // cout << i << " "<<book[i][j] << endl;
    }
  }
  rep(i, 0, 10) {
    rep(j, 0, book[i].size()) { book[i][j] += j * (j + 1); }
  }
  cout << f() << endl;
  // cout << book[1][3] << endl;
  return 0;
}