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

#include <bits/stdc++.h>
#define ll long long
#define INF 1000000005
#define MOD 1000000007
#define EPS 1e-10
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define rrep(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define srep(i, s, t) for (int i = (int)(s); i < (int)(t); ++i)
#define each(a, b) for (auto &(a) : (b))
#define all(v) (v).begin(), (v).end()
#define len(v) (int)(v).size()
#define zip(v) sort(all(v)), v.erase(unique(all(v)), v.end())
#define cmx(x, y) x = max(x, y)
#define cmn(x, y) x = min(x, y)
#define fi first
#define se second
#define pb push_back
#define show(x) cout << #x << " = " << (x) << endl
#define spair(p) cout << #p << ": " << p.fi << " " << p.se << endl
#define sar(a, n)                                                              \\
  cout << #a << ":";                                                           \\
  rep(kbrni, n) cout << " " << a[kbrni];                                       \\
  cout << endl
#define svec(v)                                                                \\
  cout << #v << ":";                                                           \\
  rep(kbrni, v.size()) cout << " " << v[kbrni];                                \\
  cout << endl
#define svecp(v)                                                               \\
  cout << #v << ":";                                                           \\
  each(kbrni, v) cout << " {" << kbrni.first << ":" << kbrni.second << "}";    \\
  cout << endl
#define sset(s)                                                                \\
  cout << #s << ":";                                                           \\
  each(kbrni, s) cout << " " << kbrni;                                         \\
  cout << endl
#define smap(m)                                                                \\
  cout << #m << ":";                                                           \\
  each(kbrni, m) cout << " {" << kbrni.first << ":" << kbrni.second << "}";    \\
  cout << endl

using namespace std;

typedef pair<int, int> P;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<double> vd;
typedef vector<P> vp;
typedef vector<string> vs;

const int MAX_N = 10005;

struct edge {
  int to, cost;
};

vector<edge> G[MAX_N];
bool ok[MAX_N];
int sub[MAX_N];
int n, t, k;

void dfs(int u, int p) {
  sub[u] = 1;
  each(e, G[u]) {
    int v = e.to;
    if (p != v) {
      dfs(v, u);
      sub[u] += sub[v];
    }
  }
}

vvi redfs(int u, int p) {
  vvi dp(2, vi(2, INF));
  int mx = 1;
  if (ok[u]) {
    dp[1][1] = 0;
  } else {
    dp[0][1] = 0;
  }
  sort(all(G[u]), [&](edge &a, edge &b) {
    return (b.to == p) ? true : sub[a.to] < sub[b.to];
  });
  each(e, G[u]) {
    cmn(mx, k + 1);
    int v = e.to, c = e.cost;
    if (p != v) {
      vvi res = redfs(v, u);
      vvi dp2(2, vi(mx + len(res[0]), INF));
      srep(i, 1, mx + 1) {
        srep(j, 1, len(res[0])) {
          if (ok[u]) {
            cmn(dp2[1][i + j - 1], dp[1][i] + res[0][j]);
            cmn(dp2[1][i + j - 1], dp[1][i] + res[1][j]);
            cmn(dp2[1][i + j], dp[1][i] + res[1][j] + c);
          } else {
            cmn(dp2[0][i + j - 1], dp[0][i] + res[0][j]);
            cmn(dp2[0][i + j], dp[0][i] + res[1][j] + c);
            cmn(dp2[1][i + j - 1], dp[0][i] + res[1][j]);
            cmn(dp2[1][i + j - 1], dp[1][i] + res[0][j]);
            cmn(dp2[1][i + j - 1], dp[1][i] + res[1][j]);
            cmn(dp2[1][i + j], dp[1][i] + res[1][j] + c);
          }
        }
      }
      mx += len(res[0]) - 1;
      dp[0].resize(mx + 1), dp[1].resize(mx + 1);
      rep(i, 2) {
        rep(j, mx + 1) { dp[i][j] = dp2[i][j]; }
      }
    }
  }
  return dp;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int id = 0;
  while (1) {
    cin >> n >> t >> k;
    if (n == 0) {
      break;
    }
    rep(i, n) { G[i].clear(), ok[i] = false; }
    rep(i, n - 1) {
      int a, b, c;
      cin >> a >> b >> c;
      G[a - 1].pb((edge){b - 1, c}), G[b - 1].pb((edge){a - 1, c});
    }
    rep(i, t) {
      int a;
      cin >> a;
      --a;
      ok[a] = true;
    }
    if (k == 0) {
      cout << "Case " << ++id << ": 0" << endl;
      continue;
    }
    dfs(0, -1);
    vvi ans = redfs(0, -1);
    cout << "Case " << ++id << ": " << ans[1][k + 1] << endl;
  }
  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
#define ll long long
#define INF 1000000005
#define MOD 1000000007
#define EPS 1e-10
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define rrep(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define srep(i, s, t) for (int i = (int)(s); i < (int)(t); ++i)
#define each(a, b) for (auto &(a) : (b))
#define all(v) (v).begin(), (v).end()
#define len(v) (int)(v).size()
#define zip(v) sort(all(v)), v.erase(unique(all(v)), v.end())
#define cmx(x, y) x = max(x, y)
#define cmn(x, y) x = min(x, y)
#define fi first
#define se second
#define pb push_back
#define show(x) cout << #x << " = " << (x) << endl
#define spair(p) cout << #p << ": " << p.fi << " " << p.se << endl
#define sar(a, n)                                                              \\
  cout << #a << ":";                                                           \\
  rep(kbrni, n) cout << " " << a[kbrni];                                       \\
  cout << endl
#define svec(v)                                                                \\
  cout << #v << ":";                                                           \\
  rep(kbrni, v.size()) cout << " " << v[kbrni];                                \\
  cout << endl
#define svecp(v)                                                               \\
  cout << #v << ":";                                                           \\
  each(kbrni, v) cout << " {" << kbrni.first << ":" << kbrni.second << "}";    \\
  cout << endl
#define sset(s)                                                                \\
  cout << #s << ":";                                                           \\
  each(kbrni, s) cout << " " << kbrni;                                         \\
  cout << endl
#define smap(m)                                                                \\
  cout << #m << ":";                                                           \\
  each(kbrni, m) cout << " {" << kbrni.first << ":" << kbrni.second << "}";    \\
  cout << endl

using namespace std;

typedef pair<int, int> P;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<double> vd;
typedef vector<P> vp;
typedef vector<string> vs;

const int MAX_N = 10005;

struct edge {
  int to, cost;
};

vector<edge> G[MAX_N];
bool ok[MAX_N];
int sub[MAX_N];
int n, t, k;

void dfs(int u, int p) {
  sub[u] = 1;
  each(e, G[u]) {
    int v = e.to;
    if (p != v) {
      dfs(v, u);
      sub[u] += sub[v];
    }
  }
}

vvi redfs(int u, int p) {
  vvi dp(2, vi(sub[u] + 1, INF));
  int mx = 1;
  if (ok[u]) {
    dp[1][1] = 0;
  } else {
    dp[0][1] = 0;
  }
  sort(all(G[u]), [&](edge &a, edge &b) {
    return (b.to == p) ? true : sub[a.to] < sub[b.to];
  });
  each(e, G[u]) {
    cmn(mx, k + 1);
    int v = e.to, c = e.cost;
    if (p != v) {
      vvi res = redfs(v, u);
      vvi dp2(2, vi(mx + len(res[0]), INF));
      srep(i, 1, mx + 1) {
        srep(j, 1, len(res[0])) {
          if (ok[u]) {
            cmn(dp2[1][i + j - 1], dp[1][i] + res[0][j]);
            cmn(dp2[1][i + j - 1], dp[1][i] + res[1][j]);
            cmn(dp2[1][i + j], dp[1][i] + res[1][j] + c);
          } else {
            cmn(dp2[0][i + j - 1], dp[0][i] + res[0][j]);
            cmn(dp2[0][i + j], dp[0][i] + res[1][j] + c);
            cmn(dp2[1][i + j - 1], dp[0][i] + res[1][j]);
            cmn(dp2[1][i + j - 1], dp[1][i] + res[0][j]);
            cmn(dp2[1][i + j - 1], dp[1][i] + res[1][j]);
            cmn(dp2[1][i + j], dp[1][i] + res[1][j] + c);
          }
        }
      }
      mx += len(res[0]) - 1;
      rep(i, 2) {
        rep(j, mx + 1) { dp[i][j] = dp2[i][j]; }
      }
    }
  }
  return dp;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int id = 0;
  while (1) {
    cin >> n >> t >> k;
    if (n == 0) {
      break;
    }
    rep(i, n) { G[i].clear(), ok[i] = false; }
    rep(i, n - 1) {
      int a, b, c;
      cin >> a >> b >> c;
      G[a - 1].pb((edge){b - 1, c}), G[b - 1].pb((edge){a - 1, c});
    }
    rep(i, t) {
      int a;
      cin >> a;
      --a;
      ok[a] = true;
    }
    if (k == 0) {
      cout << "Case " << ++id << ": 0" << endl;
      continue;
    }
    dfs(0, -1);
    vvi ans = redfs(0, -1);
    cout << "Case " << ++id << ": " << ans[1][k + 1] << endl;
  }
  return 0;
}