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

#include <bits/stdc++.h>
using namespace std;

#define for_(i, a, b) for (int i = (a); i < (b); ++i)

const int mINF = (-5) * (int)1e8;

struct Edge {
  int to, cost;
};

void maxUpdate(int &a, int b) { a = max(a, b); }

struct Pocket {
  int n, mode;
  vector<int> dp[2][2][2][2];

  Pocket(vector<int> &cost_vec, int _mode_) {
    n = cost_vec.size() + (1 - _mode_);
    mode = _mode_;

    for_(i, 0, 2) for_(j, 0, 2) for_(ii, 0, 2) for_(jj, 0, 2) dp[i][j][ii][jj]
        .assign(n + 10, mINF);
    dp[mode][0][0][0][0] = 0;

    for_(i, 0, n) {
      int ci = i & 1, ni = 1 - ci;

      for_(cnt, 0, n + 1) for_(j, 0, 2) for_(k, 0, 2) dp[mode][ni][j][k][cnt] =
          mINF;

      for_(cnt, 0, n + 1) for_(j, 0, 2) for_(k, 0, 2) {
        if (dp[mode][ci][j][k][cnt] == mINF)
          continue;

        maxUpdate(dp[mode][ni][0][k][cnt], dp[mode][ci][j][k][cnt]);

        if (i > 0 && j && cost_vec[i - 1] == 0)
          continue;
        if (mode && i == n - 1 && k && cost_vec[n - 1] == 0)
          continue;

        maxUpdate(dp[mode][ni][1][k | (i == 0)][cnt + 1],
                  dp[mode][ci][j][k][cnt] +
                      (cnt != 0 && j ? cost_vec[i - 1] : 0) +
                      (mode && i == n - 1 && k ? cost_vec[n - 1] : 0));
      }
    }
  }

  int getCost(int k) {
    int res = mINF;
    for_(i, 0, 2) for_(j, 0, 2) maxUpdate(res, dp[mode][n & 1][i][j][k]);
    return res;
  }

  void dump() {
    for_(k, 0, n + 1) cout << getCost(k) << " ";
    cout << endl;
  }
};

int N, M, K, dp[2020][2020];
vector<vector<Edge>> edges;
vector<bool> vis;

void dfs(int v, int p, vector<int> &cost_vec) {
  if (vis[v])
    return;
  vis[v] = true;

  for (Edge e : edges[v]) {
    if (e.to != p) {
      cost_vec.push_back(e.cost);
      dfs(e.to, v, cost_vec);
    }
  }
}

void solve() {
  vis.assign(N, false);

  vector<Pocket> pockets;

  for_(sz, 1, 3) {
    for_(v, 0, N) {
      if (vis[v] || (int)edges[v].size() != sz)
        continue;
      vector<int> cost_vec;
      if (sz == 2)
        edges[v].pop_back();
      dfs(v, -1, cost_vec);
      pockets.push_back(Pocket(cost_vec, sz == 2));
      /*
      int n = cost_vec.size();
      cout << v << ", " << n << " : ";
      for_(i,0,n) cout << cost_vec[i] << " ";
      cout << endl;
      pockets.back().dump();
      */
    }
  }

  int P = pockets.size(), cnt0 = 0;

  for_(v, 0, N) cnt0 += (edges[v].size() == 0);

  for_(i, 0, P + 1) for_(j, 0, K + 1) dp[i][j] = mINF;

  for_(i, 0, cnt0 + 1) dp[0][i] = 0;

  for_(i, 0, P) {
    int n = pockets[i].n;

    for_(j, 0, K + 1) {
      if (dp[i][j] == mINF)
        continue;
      for_(k, 0, n + 1)
          maxUpdate(dp[i + 1][j + k], dp[i][j] + pockets[i].getCost(k));
    }
  }

  if (dp[P][K] == mINF)
    cout << "Impossible" << endl;
  else
    cout << dp[P][K] << endl;
}

int main() {
  cin >> N >> M >> K;

  edges.assign(N, vector<Edge>());

  for_(i, 0, M) {
    int u, v, c;
    cin >> u >> v >> c;
    --u;
    --v;
    edges[u].push_back(Edge{v, c});
    edges[v].push_back(Edge{u, c});
  }

  solve();
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define for_(i, a, b) for (int i = (a); i < (b); ++i)

const int mINF = (-5) * (int)1e8;

struct Edge {
  int to, cost;
};

void maxUpdate(int &a, int b) { a = max(a, b); }

struct Pocket {
  int n, mode, dp[2][2][2020][2][2];

  Pocket(vector<int> &cost_vec, int _mode_) {
    n = cost_vec.size() + (1 - _mode_);
    mode = _mode_;

    for_(cnt, 0, n + 1) for_(j, 0, 2) for_(k, 0, 2) dp[mode][0][cnt][j][k] =
        mINF;
    dp[mode][0][0][0][0] = 0;

    for_(i, 0, n) {
      int ci = i & 1, ni = 1 - ci;

      for_(cnt, 0, n + 1) for_(j, 0, 2) for_(k, 0, 2) dp[mode][ni][cnt][j][k] =
          mINF;

      for_(cnt, 0, n + 1) for_(j, 0, 2) for_(k, 0, 2) {
        if (dp[mode][ci][cnt][j][k] == mINF)
          continue;

        maxUpdate(dp[mode][ni][cnt][0][k], dp[mode][ci][cnt][j][k]);

        if (i > 0 && j && cost_vec[i - 1] == 0)
          continue;
        if (mode && i == n - 1 && k && cost_vec[n - 1] == 0)
          continue;

        maxUpdate(dp[mode][ni][cnt + 1][1][k | (i == 0)],
                  dp[mode][ci][cnt][j][k] +
                      (cnt != 0 && j ? cost_vec[i - 1] : 0) +
                      (mode && i == n - 1 && k ? cost_vec[n - 1] : 0));
      }
    }
  }

  int getCost(int k) {
    int res = mINF;
    for_(i, 0, 2) for_(j, 0, 2) maxUpdate(res, dp[mode][n & 1][k][i][j]);
    return res;
  }

  void dump() {
    for_(k, 0, n + 1) cout << getCost(k) << " ";
    cout << endl;
  }
};

int N, M, K, dp[2020][2020];
vector<vector<Edge>> edges;
vector<bool> vis;

void dfs(int v, int p, vector<int> &cost_vec) {
  if (vis[v])
    return;
  vis[v] = true;

  for (Edge e : edges[v]) {
    if (e.to != p) {
      cost_vec.push_back(e.cost);
      dfs(e.to, v, cost_vec);
    }
  }
}

void solve() {
  vis.assign(N, false);

  vector<Pocket> pockets;

  for_(sz, 1, 3) {
    for_(v, 0, N) {
      if (vis[v] || (int)edges[v].size() != sz)
        continue;
      vector<int> cost_vec;
      if (sz == 2)
        edges[v].pop_back();
      dfs(v, -1, cost_vec);
      pockets.push_back(Pocket(cost_vec, sz == 2));
      /*
      int n = cost_vec.size();
      cout << v << ", " << n << " : ";
      for_(i,0,n) cout << cost_vec[i] << " ";
      cout << endl;
      pockets.back().dump();
      */
    }
  }

  int P = pockets.size(), cnt0 = 0;

  for_(v, 0, N) cnt0 += (edges[v].size() == 0);

  for_(i, 0, P + 1) for_(j, 0, K + 1) dp[i][j] = mINF;

  for_(i, 0, cnt0 + 1) dp[0][i] = 0;

  for_(i, 0, P) {
    int n = pockets[i].n;

    for_(j, 0, K + 1) {
      if (dp[i][j] == mINF)
        continue;
      for_(k, 0, n + 1)
          maxUpdate(dp[i + 1][j + k], dp[i][j] + pockets[i].getCost(k));
    }
  }

  if (dp[P][K] == mINF)
    cout << "Impossible" << endl;
  else
    cout << dp[P][K] << endl;
}

int main() {
  cin >> N >> M >> K;

  edges.assign(N, vector<Edge>());

  for_(i, 0, M) {
    int u, v, c;
    cin >> u >> v >> c;
    --u;
    --v;
    edges[u].push_back(Edge{v, c});
    edges[v].push_back(Edge{u, c});
  }

  solve();
}