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

// Template {{{
#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
using namespace std;
typedef long long LL;

#ifdef LOCAL
#include "contest.h"
#else
#define dump(x)
#endif

class range {
  struct Iterator {
    int val, inc;
    int operator*() { return val; }
    bool operator!=(Iterator &rhs) { return val < rhs.val; }
    void operator++() { val += inc; }
  };
  Iterator i, n;

public:
  range(int e) : i({0, 1}), n({e, 1}) {}
  range(int b, int e) : i({b, 1}), n({e, 1}) {}
  range(int b, int e, int inc) : i({b, inc}), n({e, inc}) {}
  Iterator &begin() { return i; }
  Iterator &end() { return n; }
};

const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};
inline bool valid(int x, int w) { return 0 <= x && x < w; }

void iostream_init() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.setf(ios::fixed);
  cout.precision(12);
}
//}}}

namespace {
struct Edge {
  int dst, cost;
};
typedef vector<Edge> Node;
typedef vector<Node> Graph;
Graph G;
typedef tuple<int, int, int> TP;
vector<TP> es;
vector<bool> is_base;
int N, T, K;

set<int> used[10000];

bool find(int u, int p) {
  if (is_base[u])
    return true;
  for (Edge e : G[u]) {
    if (e.dst == p)
      continue;
    if (used[u].count(e.dst))
      continue;
    if (find(e.dst, u))
      return true;
  }
  return false;
}

int solve() {
  REP(i, N) used[i].clear();
  int ans = 0;
  for (auto tp : es) {
    if (K == 0)
      break;
    int w, a, b;
    tie(w, a, b) = tp;
    if (find(a, b) && find(b, a)) {
      K--;
      ans += w;
      used[a].insert(b);
      used[b].insert(a);
    }
  }
  return ans;
}
} // namespace

int main() {
  iostream_init();
  int Case = 0;
  while (cin >> N >> T >> K && N > 0) {
    cout << "Case " << ++Case << ": ";
    G.assign(N, Node());
    es.clear();
    is_base.assign(N, false);

    REP(i, N - 1) {
      int a, b, w;
      cin >> a >> b >> w;
      a--;
      b--;
      G[a].push_back({b, w});
      G[b].push_back({a, w});
      es.push_back(TP(w, a, b));
    }
    sort(es.begin(), es.end());

    REP(i, T) {
      int base;
      cin >> base;
      base--;
      is_base[base] = true;
    }

    cout << solve() << endl;
  }
  return 0;
}

/* vim:set foldmethod=marker: */

Error: MLE

Fixed submission:

// Template {{{
#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
using namespace std;
typedef long long LL;

#ifdef LOCAL
#include "contest.h"
#else
#define dump(x)
#endif

class range {
  struct Iterator {
    int val, inc;
    int operator*() { return val; }
    bool operator!=(Iterator &rhs) { return val < rhs.val; }
    void operator++() { val += inc; }
  };
  Iterator i, n;

public:
  range(int e) : i({0, 1}), n({e, 1}) {}
  range(int b, int e) : i({b, 1}), n({e, 1}) {}
  range(int b, int e, int inc) : i({b, inc}), n({e, inc}) {}
  Iterator &begin() { return i; }
  Iterator &end() { return n; }
};

const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};
inline bool valid(int x, int w) { return 0 <= x && x < w; }

void iostream_init() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.setf(ios::fixed);
  cout.precision(12);
}
//}}}

namespace {
struct Edge {
  int dst, cost;
};
typedef vector<Edge> Node;
typedef vector<Node> Graph;
Graph G;
vector<bool> is_base;
int N, T, K;

int dp[10001][10001][2];
int next[10001][2];
int size[10001];
void calc(int u, int p) {
  for (Edge e : G[u]) {
    if (e.dst != p) {
      calc(e.dst, u);
    }
  }

  size[u] = 1;
  for (Edge e : G[u]) {
    if (e.dst != p) {
      size[u] += size[e.dst];
    }
  }

  dp[u][0][is_base[u]] = 0;
  for (Edge e : G[u]) {
    int c = e.dst;
    if (c == p)
      continue;

    // memset(next, 0x3f, sizeof(next));
    REP(i, K + 1) REP(b, 2) next[i][b] = 1e9;

    for (int i = 0; i <= K; i++) {
      for (int j = 0; j < size[c] && i + j <= K; j++) {
        // connect
        {
          const int ni = i + j;
          next[ni][0] = min(next[ni][0], dp[u][i][0] + dp[c][j][0]);
          next[ni][1] = min(next[ni][1], dp[u][i][1] + dp[c][j][0]);
          next[ni][1] = min(next[ni][1], dp[u][i][0] + dp[c][j][1]);
          next[ni][1] = min(next[ni][1], dp[u][i][1] + dp[c][j][1]);
        }
        // cut
        {
          const int ni = i + j + 1;
          if (ni <= K) {
            next[ni][0] = min(next[ni][0], dp[u][i][0] + dp[c][j][1] + e.cost);
            next[ni][1] = min(next[ni][1], dp[u][i][1] + dp[c][j][1] + e.cost);
          }
        }
      }
    }

    // memcpy(dp[u], next, sizeof(next));
    REP(i, K + 1) REP(b, 2) dp[u][i][b] = next[i][b];
  }
}
void output(int u, int k, int b) {
  printf("dp[%d][%d][%d] = %d\\n", u, k, b, dp[u][k][b]);
}
} // namespace

int main() {
  iostream_init();
  int Case = 0;
  while (cin >> N >> T >> K && N > 0) {
    cout << "Case " << ++Case << ": ";
    G.assign(N, Node());
    is_base.assign(N, false);

    memset(dp, 0x3f, sizeof(dp));

    REP(i, N - 1) {
      int a, b, w;
      cin >> a >> b >> w;
      a--;
      b--;
      G[a].push_back({b, w});
      G[b].push_back({a, w});
    }

    REP(i, T) {
      int base;
      cin >> base;
      base--;
      is_base[base] = true;
    }

    calc(0, -1);
    cout << dp[0][K][1] << endl;
  }
  return 0;
}

/* vim:set foldmethod=marker: */