Problem p03144 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cstdio>
#include <map>
#include <set>
#include <string>
#include <vector>

using int64 = long long;
const int64 inf = 1000000000000ll;
struct Color {
  std::vector<int> costs;
  size_t pos{0};
  int id;

  int64 cost_n(int n) {
    if (n == 0) {
      if (pos < 3) {
        return -inf;
      }
      return -costs[pos - 1];
    }
    if (pos == 0 && n < 2) {
      return -1;
    }
    if (pos + n > costs.size()) {
      return -1;
    }
    int64 res = 0;
    for (int i = 0; i < n; i++) {
      res += costs[pos + i];
    }
    return res;
  }
  auto key(int d) { return std::make_pair(-cost_n(d), id); }
};
void relax(int64 &x, int64 v) {
  if (x < v) {
    x = v;
  }
}

std::vector<int64> solve(int n, int k, std::vector<Color> colors) {
  for (auto &c : colors) {
    std::sort(c.costs.begin(), c.costs.end(), std::greater<>());
  }
  std::vector<int64> res(n + 1, -1);
  const int maxd = 3;
  std::set<std::pair<int64, int>> best[maxd + 1];
  auto add = [&](int id) {
    for (int d = 0; d <= maxd; d++) {
      best[d].insert(colors[id].key(d));
    }
  };
  auto remove = [&](int id) {
    for (int d = 0; d <= maxd; d++) {
      best[d].erase(colors[id].key(d));
    }
  };
  for (int i = 0; i < k; i++) {
    add(i);
  }
  int cnt = 0;
  int64 cost = 0;
  while (cnt < n) {
    for (int d = 1; d <= maxd; d++) {
      auto x = *best[d].begin();
      if (x.first > 0) {
        continue;
      }
      relax(res[cnt + d], cost - x.first);
    }
    auto one = *best[1].begin();
    auto two = *best[2].begin();
    if (one.first > 0 || one.first * 2 >= two.first) {
      cnt += 2;
      cost -= two.first;
      remove(two.second);
      colors[two.second].pos += 2;
      add(two.second);
      // printf("two");
    } else {
      cnt += 1;
      cost -= one.first;
      remove(one.second);
      colors[one.second].pos += 1;
      add(one.second);
      // printf("one");
    }

    // printf("?%d %lld\n", cnt, cost);
    relax(res[cnt], cost);

    {
      auto x = *best[0].begin();
      if (x.first != inf) {
        relax(res[cnt - 1], cost - x.first);
      }
    }
  }
  res[0] = 0;
  return res;
}

void dfs(int i, int cnt, int64 cost, std::vector<int64> &res,
         std::vector<Color> &colors, int k) {
  if (i == k) {
    relax(res[cnt], cost);
    return;
  }
  int add_cnt = 0;
  int64 add_cost = 0;
  dfs(i + 1, cnt, cost, res, colors, k);
  for (int x : colors[i].costs) {
    add_cnt++;
    add_cost += x;
    if (add_cnt >= 2) {
      dfs(i + 1, cnt + add_cnt, cost + add_cost, res, colors, k);
    }
  }
}
std::vector<int64> solve_slow(int n, int k, std::vector<Color> colors) {
  for (auto &c : colors) {
    std::sort(c.costs.begin(), c.costs.end(), std::greater<>());
  }
  std::vector<int64> res(n + 1, -1);
  dfs(0, 0, 0, res, colors, k);
  return res;
}
void stress() {
  int test_i = 0;
  while (true) {
    printf("test %d\n", test_i++);
    int n = 0;
    int k = rand() % 2 + 5;
    std::vector<Color> colors(k);
    for (int i = 0; i < k; i++) {
      int cnt = rand() % 4 + 2;
      colors[i].id = i;
      for (int j = 0; j < cnt; j++) {
        colors[i].costs.push_back(rand() % 5 + 1);
        n++;
      }
    }
    for (auto &c : colors) {
      std::sort(c.costs.begin(), c.costs.end(), std::greater<>());
    }
    auto res = solve(n, k, colors);
    auto res2 = solve_slow(n, k, colors);
    if (res != res2) {
      printf("mismatch\n");
      printf("k = %d\n", k);
      for (int i = 0; i < k; i++) {
        for (int c : colors[i].costs) {
          printf("%d ", c);
        }
        printf("\n");
      }
      for (size_t i = 2; i < res.size(); i++) {
        printf("%d: %lld vs %lld\n", (int)i, res[i], res2[i]);
      }
      std::exit(1);
    }
  }
}

int main(void) {
  int n, k;
  // stress();
  while (scanf("%d%d", &n, &k) == 2) {
    std::vector<Color> colors(k);
    for (int i = 0; i < n; i++) {
      int c, v;
      scanf("%d%d", &c, &v);
      colors[c - 1].costs.push_back(v);
    }
    for (int i = 0; i < k; i++) {
      colors[i].id = i;
    }
    auto res = solve(n, k, colors);
    auto res2 = solve_slow(n, k, colors);
    if (res != res2) {
      printf("mismatch\n");
    }
    for (int i = 1; i <= n; i++) {
      printf("%lld\n", res[i]);
    }
  }

  return 0;
}


Fixed submission:
#include <algorithm>
#include <cstdio>
#include <map>
#include <set>
#include <string>
#include <vector>

using int64 = long long;
const int64 inf = 1000000000000ll;
struct Color {
  std::vector<int> costs;
  size_t pos{0};
  int id;

  int64 cost_n(int n) {
    if (n == 0) {
      if (pos < 3) {
        return -inf;
      }
      return -costs[pos - 1];
    }
    if (pos == 0 && n < 2) {
      return -1;
    }
    if (pos + n > costs.size()) {
      return -1;
    }
    int64 res = 0;
    for (int i = 0; i < n; i++) {
      res += costs[pos + i];
    }
    return res;
  }
  auto key(int d) { return std::make_pair(-cost_n(d), id); }
};
void relax(int64 &x, int64 v) {
  if (x < v) {
    x = v;
  }
}

std::vector<int64> solve(int n, int k, std::vector<Color> colors) {
  for (auto &c : colors) {
    std::sort(c.costs.begin(), c.costs.end(), std::greater<>());
  }
  std::vector<int64> res(n + 1, -1);
  const int maxd = 3;
  std::set<std::pair<int64, int>> best[maxd + 1];
  auto add = [&](int id) {
    for (int d = 0; d <= maxd; d++) {
      best[d].insert(colors[id].key(d));
    }
  };
  auto remove = [&](int id) {
    for (int d = 0; d <= maxd; d++) {
      best[d].erase(colors[id].key(d));
    }
  };
  for (int i = 0; i < k; i++) {
    add(i);
  }
  int cnt = 0;
  int64 cost = 0;
  while (cnt < n) {
    for (int d = 1; d <= maxd; d++) {
      auto x = *best[d].begin();
      if (x.first > 0) {
        continue;
      }
      relax(res[cnt + d], cost - x.first);
    }
    auto one = *best[1].begin();
    auto two = *best[2].begin();
    if (one.first > 0 || one.first * 2 >= two.first) {
      cnt += 2;
      cost -= two.first;
      remove(two.second);
      colors[two.second].pos += 2;
      add(two.second);
      // printf("two");
    } else {
      cnt += 1;
      cost -= one.first;
      remove(one.second);
      colors[one.second].pos += 1;
      add(one.second);
      // printf("one");
    }

    // printf("?%d %lld\n", cnt, cost);
    relax(res[cnt], cost);

    {
      auto x = *best[0].begin();
      if (x.first != inf) {
        relax(res[cnt - 1], cost - x.first);
      }
    }
  }
  res[0] = 0;
  return res;
}

void dfs(int i, int cnt, int64 cost, std::vector<int64> &res,
         std::vector<Color> &colors, int k) {
  if (i == k) {
    relax(res[cnt], cost);
    return;
  }
  int add_cnt = 0;
  int64 add_cost = 0;
  dfs(i + 1, cnt, cost, res, colors, k);
  for (int x : colors[i].costs) {
    add_cnt++;
    add_cost += x;
    if (add_cnt >= 2) {
      dfs(i + 1, cnt + add_cnt, cost + add_cost, res, colors, k);
    }
  }
}
std::vector<int64> solve_slow(int n, int k, std::vector<Color> colors) {
  for (auto &c : colors) {
    std::sort(c.costs.begin(), c.costs.end(), std::greater<>());
  }
  std::vector<int64> res(n + 1, -1);
  dfs(0, 0, 0, res, colors, k);
  return res;
}
void stress() {
  int test_i = 0;
  while (true) {
    printf("test %d\n", test_i++);
    int n = 0;
    int k = rand() % 2 + 5;
    std::vector<Color> colors(k);
    for (int i = 0; i < k; i++) {
      int cnt = rand() % 4 + 2;
      colors[i].id = i;
      for (int j = 0; j < cnt; j++) {
        colors[i].costs.push_back(rand() % 5 + 1);
        n++;
      }
    }
    for (auto &c : colors) {
      std::sort(c.costs.begin(), c.costs.end(), std::greater<>());
    }
    auto res = solve(n, k, colors);
    auto res2 = solve_slow(n, k, colors);
    if (res != res2) {
      printf("mismatch\n");
      printf("k = %d\n", k);
      for (int i = 0; i < k; i++) {
        for (int c : colors[i].costs) {
          printf("%d ", c);
        }
        printf("\n");
      }
      for (size_t i = 2; i < res.size(); i++) {
        printf("%d: %lld vs %lld\n", (int)i, res[i], res2[i]);
      }
      std::exit(1);
    }
  }
}

int main(void) {
  int n, k;
  // stress();
  while (scanf("%d%d", &n, &k) == 2) {
    std::vector<Color> colors(k);
    for (int i = 0; i < n; i++) {
      int c, v;
      scanf("%d%d", &c, &v);
      colors[c - 1].costs.push_back(v);
    }
    for (int i = 0; i < k; i++) {
      colors[i].id = i;
    }
    auto res = solve(n, k, colors);
    // auto res2 = solve_slow(n, k, colors);
    // if (res != res2) {
    // printf("mismatch\n");
    // }
    for (int i = 1; i <= n; i++) {
      printf("%lld\n", res[i]);
    }
  }

  return 0;
}