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

#include <bits/stdc++.h>
#define INF 1000000007
using namespace std;
typedef long long ll;
typedef pair<int, int> P;

struct diameter {
  static const int MAX = 10000;
  vector<int> G[MAX];
  int n;
  int d[MAX];
  bool vis[MAX];
  int d2[2][MAX];

  void bfs(int s) {
    fill(d, d + n, INF);
    queue<int> que;
    que.push(s);
    d[s] = 0;
    int v;
    while (que.size()) {
      v = que.front();
      que.pop();
      for (int i = 0; i < G[v].size(); i++) {
        int e = G[v][i];
        if (d[e] == INF) {
          d[e] = d[v] + 1;
          que.push(e);
        }
      }
    }
  }
  P diameter_points(int s = 0) {
    bfs(s);
    int maxv = 0;
    int tgt = 0;
    for (int i = 0; i < n; i++) {
      if (d[i] == INF)
        continue;
      if (maxv < d[i]) {
        maxv = d[i];
        tgt = i;
      }
    }
    bfs(tgt);
    maxv = 0;
    int v = 0;
    for (int i = 0; i < n; i++) {
      if (d[i] == INF)
        continue;
      if (maxv < d[i]) {
        maxv = d[i];
        v = i;
      }
    }
    return P(tgt, v);
  }
  vector<int> middle_points() {
    P p = diameter_points();
    bfs(p.first);
    for (int i = 0; i < n; i++) {
      d2[0][i] = d[i];
    }
    bfs(p.second);
    for (int i = 0; i < n; i++) {
      d2[1][i] = d[i];
    }
    vector<int> cid;
    if (d2[0][p.second] % 2 == 0) {
      for (int i = 0; i < n; i++) {
        if (d2[0][p.second] == d2[0][i] * 2 && d2[1][p.first] == d2[1][i] * 2) {
          cid.push_back(i);
        }
      }
    } else {
      for (int i = 0; i < n; i++) {
        if (d2[0][p.second] == (d2[0][i] + d2[1][i]) &&
            abs(d2[0][i] - d2[1][i]) == 1) {
          cid.push_back(i);
        }
      }
    }
    return cid;
  }
};

struct edge {
  int t, c;
  edge() {}
  edge(int tt, int cc) {
    t = tt;
    c = cc;
  }
};

diameter di;
vector<edge> G[10001];
int n, t, k;
bool flag[10001];
int siz[10001];
int *dp[10001][2];
int tmp[2][10001];

int dfs(int v, int p) {
  int ans = 1;
  for (int i = 0; i < G[v].size(); i++) {
    int nv = G[v][i].t;
    if (nv != p) {
      ans += dfs(nv, v);
    }
  }
  return (siz[v] = ans);
}

void func(int v, int p) {
  free(dp[v][0]);
  free(dp[v][1]);
  dp[v][0] = new int[siz[v] + 1];
  dp[v][1] = new int[siz[v] + 1];
  if (flag[v]) {
    dp[v][0][1] = INF;
    dp[v][1][1] = 0;
  } else {
    dp[v][1][1] = INF;
    dp[v][0][1] = 0;
  }
  int ns = 1;
  for (int i = 0; i < G[v].size(); i++) {
    int nv = G[v][i].t;
    int nc = G[v][i].c;
    if (nv != p) {
      func(nv, v);
      for (int j = 0; j <= ns + siz[nv]; j++) {
        tmp[0][j] = INF;
        tmp[1][j] = INF;
      }
      for (int i = 1; i <= ns; i++) {
        for (int j = 1; j <= siz[nv]; j++) {
          int l = i + j;
          tmp[0][l] = min(tmp[0][i + j], dp[v][0][i] + dp[nv][1][j] + nc);
          tmp[1][l] = min(tmp[1][i + j], dp[v][1][i] + dp[nv][1][j] + nc);
          tmp[0][l - 1] = min(tmp[0][i + j - 1], dp[v][0][i] + dp[nv][0][j]);
          tmp[1][l - 1] = min(tmp[1][i + j - 1], dp[v][1][i] + dp[nv][0][j]);
          tmp[1][l - 1] = min(tmp[1][i + j - 1], dp[v][0][i] + dp[nv][1][j]);
          tmp[1][l - 1] = min(tmp[1][i + j - 1], dp[v][1][i] + dp[nv][1][j]);
        }
      }
      for (int j = 1; j <= ns + siz[nv]; j++) {
        dp[v][0][j] = tmp[0][j];
        dp[v][1][j] = tmp[1][j];
      }
      ns += siz[nv];
    }
  }
}

void solve() {
  int root = di.middle_points()[0];
  dfs(root, -1);
  func(root, -1);
  int ans = INF;
  for (int i = k + 1; i <= n; i++) {
    ans = min(ans, dp[root][1][i]);
  }
  printf("%d\\n", ans);
}

int main(void) {
  int casenum = 1;
  while (1) {
    scanf("%d%d%d", &n, &t, &k);
    if (n == 0)
      break;
    di.n = n;
    memset(flag, false, sizeof(flag));
    memset(siz, false, sizeof(siz));
    for (int i = 0; i < n; i++) {
      G[i].clear();
      di.G[i].clear();
    }
    for (int i = 0; i < n - 1; i++) {
      int a, b, c;
      scanf("%d%d%d", &a, &b, &c);
      a--;
      b--;
      G[a].push_back(edge(b, c));
      G[b].push_back(edge(a, c));
      di.G[a].push_back(b);
      di.G[b].push_back(a);
    }
    for (int i = 0; i < t; i++) {
      int a;
      scanf("%d", &a);
      a--;
      flag[a] = true;
    }
    printf("Case %d: ", casenum++);
    solve();
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define INF 1000000007
using namespace std;
typedef long long ll;
typedef pair<int, int> P;

struct diameter {
  static const int MAX = 10000;
  vector<int> G[MAX];
  int n;
  int d[MAX];
  bool vis[MAX];
  int d2[2][MAX];

  void bfs(int s) {
    fill(d, d + n, INF);
    queue<int> que;
    que.push(s);
    d[s] = 0;
    int v;
    while (que.size()) {
      v = que.front();
      que.pop();
      for (int i = 0; i < G[v].size(); i++) {
        int e = G[v][i];
        if (d[e] == INF) {
          d[e] = d[v] + 1;
          que.push(e);
        }
      }
    }
  }
  P diameter_points(int s = 0) {
    bfs(s);
    int maxv = 0;
    int tgt = 0;
    for (int i = 0; i < n; i++) {
      if (d[i] == INF)
        continue;
      if (maxv < d[i]) {
        maxv = d[i];
        tgt = i;
      }
    }
    bfs(tgt);
    maxv = 0;
    int v = 0;
    for (int i = 0; i < n; i++) {
      if (d[i] == INF)
        continue;
      if (maxv < d[i]) {
        maxv = d[i];
        v = i;
      }
    }
    return P(tgt, v);
  }
  vector<int> middle_points() {
    P p = diameter_points();
    bfs(p.first);
    for (int i = 0; i < n; i++) {
      d2[0][i] = d[i];
    }
    bfs(p.second);
    for (int i = 0; i < n; i++) {
      d2[1][i] = d[i];
    }
    vector<int> cid;
    if (d2[0][p.second] % 2 == 0) {
      for (int i = 0; i < n; i++) {
        if (d2[0][p.second] == d2[0][i] * 2 && d2[1][p.first] == d2[1][i] * 2) {
          cid.push_back(i);
        }
      }
    } else {
      for (int i = 0; i < n; i++) {
        if (d2[0][p.second] == (d2[0][i] + d2[1][i]) &&
            abs(d2[0][i] - d2[1][i]) == 1) {
          cid.push_back(i);
        }
      }
    }
    return cid;
  }
};

struct edge {
  int t, c;
  edge() {}
  edge(int tt, int cc) {
    t = tt;
    c = cc;
  }
};

diameter di;
vector<edge> G[10001];
int n, t, k;
bool flag[10001];
int siz[10001];
int *dp[10001][2];
int tmp[2][10001];

int dfs(int v, int p) {
  int ans = 1;
  for (int i = 0; i < G[v].size(); i++) {
    int nv = G[v][i].t;
    if (nv != p) {
      ans += dfs(nv, v);
    }
  }
  return (siz[v] = ans);
}

void func(int v, int p) {
  free(dp[v][0]);
  free(dp[v][1]);
  dp[v][0] = new int[siz[v] + 1];
  dp[v][1] = new int[siz[v] + 1];
  if (flag[v]) {
    dp[v][0][1] = INF;
    dp[v][1][1] = 0;
  } else {
    dp[v][1][1] = INF;
    dp[v][0][1] = 0;
  }
  int ns = 1;
  for (int i = 0; i < G[v].size(); i++) {
    int nv = G[v][i].t;
    int nc = G[v][i].c;
    if (nv != p) {
      func(nv, v);
      for (int j = 0; j <= ns + siz[nv]; j++) {
        tmp[0][j] = INF;
        tmp[1][j] = INF;
      }
      for (int i = 1; i <= ns; i++) {
        for (int j = 0; j <= siz[nv]; j++) {
          tmp[0][i + j] = min(tmp[0][i + j], dp[v][0][i] + dp[nv][1][j] + nc);
          tmp[1][i + j] = min(tmp[1][i + j], dp[v][1][i] + dp[nv][1][j] + nc);
          tmp[0][i + j - 1] =
              min(tmp[0][i + j - 1], dp[v][0][i] + dp[nv][0][j]);
          tmp[1][i + j - 1] =
              min(tmp[1][i + j - 1], dp[v][1][i] + dp[nv][0][j]);
          tmp[1][i + j - 1] =
              min(tmp[1][i + j - 1], dp[v][0][i] + dp[nv][1][j]);
          tmp[1][i + j - 1] =
              min(tmp[1][i + j - 1], dp[v][1][i] + dp[nv][1][j]);
        }
      }
      for (int j = 0; j <= ns + siz[nv]; j++) {
        dp[v][0][j] = tmp[0][j];
        dp[v][1][j] = tmp[1][j];
      }
      ns += siz[nv];
    }
  }
}

void solve() {
  int root = di.middle_points()[0];
  dfs(root, -1);
  func(root, -1);
  int ans = INF;
  for (int i = k + 1; i <= n; i++) {
    ans = min(ans, dp[root][1][i]);
  }
  printf("%d\\n", ans);
}

int main(void) {
  int casenum = 1;
  while (1) {
    scanf("%d%d%d", &n, &t, &k);
    if (n == 0)
      break;
    di.n = n;
    memset(flag, false, sizeof(flag));
    memset(siz, false, sizeof(siz));
    for (int i = 0; i < n; i++) {
      G[i].clear();
      di.G[i].clear();
    }
    for (int i = 0; i < n - 1; i++) {
      int a, b, c;
      scanf("%d%d%d", &a, &b, &c);
      a--;
      b--;
      G[a].push_back(edge(b, c));
      G[b].push_back(edge(a, c));
      di.G[a].push_back(b);
      di.G[b].push_back(a);
    }
    for (int i = 0; i < t; i++) {
      int a;
      scanf("%d", &a);
      a--;
      flag[a] = true;
    }
    printf("Case %d: ", casenum++);
    solve();
  }
  return 0;
}