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

#include <algorithm>
#include <cstdio>

#define REP(i, s, n) for (int i = s; i < n; i++)
#define rep(i, n) REP(i, 0, n)

using namespace std;

struct Edge {
  int src, dst, weight;
  bool operator<(const Edge &e) const { return weight < e.weight; }
};

const int MAX = 10010;
int par[MAX];
// vector<Edge> edges;
Edge edges[MAX];

void init(int _size = MAX) { rep(i, _size) par[i] = i; }

int find(int x) {
  if (par[x] == x)
    return x;
  return par[x] = find(par[x]);
}

void unit(int x, int y) {
  x = find(x), y = find(y);
  if (x != y)
    par[x] = y;
}

bool buffer[MAX];
bool choose[MAX];
int counter[MAX];
// vector<int> bases;
int bases[MAX];
int V, t, k;

bool isValid(int index) {
  choose[index] = false;
  // map<int,int> mp;
  rep(i, V) counter[i] = 0, par[i] = i;

  rep(i, V - 1) if (choose[i]) unit(edges[i].src, edges[i].dst);
  rep(i, t)++ counter[find(bases[i])];
  rep(i, V) if (counter[find(i)] == 0) return false;
  return true;
}

void compute() {
  sort(edges, edges + V - 1);
  if (V == t) {
    int sum = 0;
    rep(i, k) sum += edges[i].weight;
    printf("%d\\n", sum);
    return;
  }
  rep(i, V) buffer[i] = true;
  int answer = 0;
  int remain = k;
  int cur = 0;
  while (cur < V - 1 && remain) {
    rep(i, V) choose[i] = buffer[i];
    if (isValid(cur)) {
      --remain;
      answer += edges[cur].weight;
      buffer[cur] = false;
    }
    ++cur;
  }
  printf("%d\\n", answer);
}

int main() {
  int CNT = 1;
  while (scanf(" %d %d %d", &V, &t, &k), V | t | k) {

    rep(i, V - 1) {
      int s, t, c;
      scanf(" %d %d %d", &s, &t, &c);
      --s, --t;
      // edges.push_back((Edge){s,t,c});
      edges[i] = (Edge){s, t, c};
    }
    rep(i, t) {
      int base;
      scanf(" %d", &base);
      --base;
      // bases.push_back(base);
      bases[i] = base;
    }
    printf("Case %d: ", CNT++);
    compute();
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdio>

#define REP(i, s, n) for (int i = s; i < n; i++)
#define rep(i, n) REP(i, 0, n)

using namespace std;

struct Edge {
  int src, dst, weight;
  bool operator<(const Edge &e) const { return weight < e.weight; }
};

const int MAX = 10010;
int par[MAX];
// vector<Edge> edges;
Edge edges[MAX];

void init(int _size = MAX) { rep(i, _size) par[i] = i; }

int find(int x) {
  if (par[x] == x)
    return x;
  return par[x] = find(par[x]);
}

void unit(int x, int y) {
  x = find(x), y = find(y);
  if (x != y)
    par[x] = y;
}

bool buffer[MAX];
bool choose[MAX];
int counter[MAX];
// vector<int> bases;
int bases[MAX];
int V, t, k;

bool isValid(int index) {
  choose[index] = false;
  // map<int,int> mp;
  rep(i, V) counter[i] = 0, par[i] = i;

  rep(i, V - 1) if (choose[i]) unit(edges[i].src, edges[i].dst);
  rep(i, t)++ counter[find(bases[i])];
  rep(i, V) if (counter[find(i)] == 0) return false;
  return true;
}

void compute() {
  sort(edges, edges + V - 1);
  rep(i, V) buffer[i] = true;
  int answer = 0;
  int remain = k;
  int cur = 0;
  while (cur < V - 1 && remain) {
    rep(i, V) choose[i] = buffer[i];
    if (isValid(cur)) {
      --remain;
      answer += edges[cur].weight;
      buffer[cur] = false;
    }
    ++cur;
  }
  printf("%d\\n", answer);
}

int main() {
  int CNT = 1;
  while (scanf(" %d %d %d", &V, &t, &k), V | t | k) {
    rep(i, V - 1) {
      int s, t, c;
      scanf(" %d %d %d", &s, &t, &c);
      --s, --t;
      // edges.push_back((Edge){s,t,c});
      edges[i] = (Edge){s, t, c};
    }
    rep(i, t) {
      int base;
      scanf(" %d", &base);
      --base;
      // bases.push_back(base);
      bases[i] = base;
    }
    printf("Case %d: ", CNT++);
    compute();
  }
  return 0;
}