leetcode-go/written-exam/kickstart-round-C/no3/main.cc

#include <algorithm>
#include <cstdio>
#include <queue>
#include <set>

using std::sort; using std::priority_queue; using std::set;

const int N = 1000;
const int A = 1000;

struct Dog {
  int p;
  int a;
  int n;
  Dog(int p = 0, int a = 0, int n = 1) : p(p), a(a), n(n) {}
  bool operator<(const Dog &that) const {
    if (this->p == that.p) return this->a < that.a;
    return this->p < that.p;
  }
} dogs[N];

struct Node {
  int t;
  int i;
  int a;
  int k;
  Node(int t, int i, int a, int k) : t(t), i(i), a(a), k(k) {}
  bool operator<(const Node &that) const {
    return that.t < this->t;
  }
};

int main() {
  int T = 0, n, k;
  scanf("%d", &T);
  for (int t = 1; t <= T; t++) {
    scanf("%d %d", &n, &k);
    for (int i = 0; i < n; i++)
      scanf("%d", &dogs[i].p);
    set<int> colors;
    for (int i = 0; i < n; i++) {
      scanf("%d", &dogs[i].a);
      colors.insert(dogs[i].a);
    }
    sort(dogs, dogs + n);
    int d = 0, aa = dogs[0].a, pp = dogs[0].p;
    for (int i = 1; i < n; i++) {
      if (dogs[i].a == aa && dogs[i].p == pp) {
        dogs[d].n++;
      } else {
        dogs[++d] = dogs[i];
        aa = dogs[d].a;
        pp = dogs[d].p;
      }
    }
    n = d + 1;
    priority_queue<Node> pq;
    pq.emplace(0, -1, *colors.begin(), 0);
    Node tp(0, -1, 0, 0);
    while (!pq.empty()) {
      tp = pq.top();
      pq.pop();
      if (k <= tp.k) break;
      int pos = tp.i == -1 ? 0 : dogs[tp.i].p;
      set<int>::iterator it = ++colors.find(tp.a);
      if (it != colors.end()) pq.emplace(tp.t + pos, -1, *it, tp.k);
      int nxt = tp.i + 1;
      for (; nxt < n && dogs[nxt].a != tp.a; nxt++) {}
      if (nxt < n) pq.emplace(tp.t + dogs[nxt].p - pos, nxt, tp.a, tp.k + dogs[nxt].n);
    }
    printf("Case #%d: %d\n", t, tp.t);
  }
  return 0;
}