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

#include <cassert>
#include <iostream>
#include <vector>

using namespace std;
using ll = long long;
constexpr ll INF = 1LL << 50;
ll s;

struct Cube {
  ll x;
  ll y;
  ll z;
};
inline ll connected(const Cube &c1, const Cube &c2) {
  const ll dx = abs(c1.x - c2.x);
  const ll dy = abs(c1.y - c2.y);
  const ll dz = abs(c1.z - c2.z);
  if (dx < s and dy < s and dz < s) {
    return ((s - dx) * (s - dy) + (s - dy) * (s - dz) + (s - dz) * (s - dx)) *
           2;
  }
  return 0;
}

struct Edge {
  int from;
  int to;
  ll cost;
};
struct Graph {
  Graph(const int n) { edge.resize(n); }
  void addEdge(const int from, const int to, const ll cost) {
    edge[from].push_back(Edge{from, to, cost});
    edge[to].push_back(Edge{to, from, cost});
  }
  vector<vector<Edge>> edge;
};

int getEdge(const Graph &g, const int s, vector<bool> &visited) {
  visited[s] = true;
  int child = 0;
  for (const auto &e : g.edge[s]) {
    if (not visited[e.to]) {
      return getEdge(g, e.to, visited);
    }
  }
  if (child == 0) {
    return s;
  }
}
struct Chain {
  bool ring = false;
  vector<ll> edges;
};
ostream &operator<<(ostream &os, const Chain &ch) {
  os << "Ring: " << boolalpha << ch.ring << "\\nEdge: {";
  for (int i = 0; i < ch.edges.size(); i++) {
    if (i != 0) {
      os << ",";
    }
    os << ch.edges[i];
  }
  os << "}" << endl;
  return os;
}
void dfs(const Graph &g, const int s, const int dep, vector<bool> &visited,
         Chain &chain) {
  visited[s] = true;
  for (const auto &e : g.edge[s]) {
    if (not visited[e.to]) {
      chain.edges.push_back(e.cost);
      dfs(g, e.to, dep, visited, chain);
    }
  }
}

int main() {
  while (true) {
    int n, k;
    cin >> n >> k >> s;
    if (n == 0 and k == 0 and s == 0) {
      break;
    }

    const ll whole = 6 * s * s * k;
    vector<Cube> cube(n);
    for (int i = 0; i < n; i++) {
      cin >> cube[i].x >> cube[i].y >> cube[i].z;
    }
    if (k == 1) {
      cout << whole << endl;
      continue;
    }

    Graph g(n);
    for (int i = 0; i < n; i++) {
      for (int j = i + 1; j < n; j++) {
        const ll cost = connected(cube[i], cube[j]);
        if (cost > 0) {
          g.addEdge(i, j, cost);
        }
      }
    }
    vector<bool> visited(n, false);
    vector<Chain> chains;
    for (int i = 0; i < n; i++) {
      if (not visited[i]) {
        vector<bool> visited2(n, false);
        const int dep = getEdge(g, i, visited2);
        Chain chain;
        chain.ring = g.edge[dep].size() == 2;
        if (chain.ring) {
          chain.edges.push_back(g.edge[dep][1].cost);
        }
        dfs(g, dep, dep, visited, chain);
        chains.push_back(chain);
      }
    }
    ll mini = INF;
    for (const auto chain : chains) {
      const bool ring = chain.ring;
      const vector<ll> &costs = chain.edges;
      if (costs.size() < k - 1) {
        continue;
      }
      if (ring) {
        ll sum = 0;
        for (int i = 0; i < k - 1; i++) {
          sum += costs[i];
        }
        if (costs.size() == k) {
          sum += costs[k - 1];
          mini = min(mini, whole - sum);
        } else {
          mini = min(mini, whole - sum);
          int head = 0;
          int tail = k - 1;
          for (int j = 0; j < costs.size();
               j++, head++, tail++, tail = tail % costs.size()) {
            sum += (costs[tail] - costs[head]);
            mini = min(mini, whole - sum);
          }
        }
      } else {
        ll sum = 0;
        for (int i = 0; i < k - 1; i++) {
          sum += costs[i];
        }
        mini = min(mini, whole - sum);
        int head = 0;
        int tail = k - 1;
        for (; tail < costs.size(); head++, tail++) {
          sum += (costs[tail] - costs[head]);
          mini = min(mini, whole - sum);
        }
      }
    }
    cout << ((mini == INF) ? -1 : mini) << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <cassert>
#include <iostream>
#include <vector>

using namespace std;
using ll = long long;
constexpr ll INF = 1LL << 50;
ll s;

struct Cube {
  ll x;
  ll y;
  ll z;
};
ll connected(const Cube &c1, const Cube &c2) {
  const ll dx = abs(c1.x - c2.x);
  const ll dy = abs(c1.y - c2.y);
  const ll dz = abs(c1.z - c2.z);
  if (dx < s and dy < s and dz < s) {
    return ((s - dx) * (s - dy) + (s - dy) * (s - dz) + (s - dz) * (s - dx)) *
           2;
  }
  return 0;
}

struct Edge {
  int from;
  int to;
  ll cost;
};
struct Graph {
  Graph(const int n) { edge.resize(n); }
  void addEdge(const int from, const int to, const ll cost) {
    edge[from].push_back(Edge{from, to, cost});
    edge[to].push_back(Edge{to, from, cost});
  }
  vector<vector<Edge>> edge;
};

int getEdge(const Graph &g, const int s, vector<bool> &visited) {
  visited[s] = true;
  int child = 0;
  for (const auto &e : g.edge[s]) {
    if (not visited[e.to]) {
      return getEdge(g, e.to, visited);
    }
  }
  if (child == 0) {
    return s;
  }
}
struct Chain {
  bool ring = false;
  vector<ll> edges;
};
ostream &operator<<(ostream &os, const Chain &ch) {
  os << "Ring: " << boolalpha << ch.ring << "\\nEdge: {";
  for (int i = 0; i < ch.edges.size(); i++) {
    if (i != 0) {
      os << ",";
    }
    os << ch.edges[i];
  }
  os << "}" << endl;
}
void dfs(const Graph &g, const int s, const int dep, vector<bool> &visited,
         Chain &chain) {
  visited[s] = true;
  for (const auto &e : g.edge[s]) {
    if (e.to == dep and chain.edges.size() >= 2) {
      if (not chain.ring) {
        chain.edges.push_back(e.cost);
        chain.ring = true;
      }
      return;
    }
    if (not visited[e.to]) {
      chain.edges.push_back(e.cost);
      dfs(g, e.to, dep, visited, chain);
    }
  }
}

int main() {
  while (true) {
    int n, k;
    cin >> n >> k >> s;
    if (n == 0 and k == 0 and s == 0) {
      break;
    }

    const ll whole = 6 * s * s * k;
    if (k == 1) {
      cout << whole << endl;
      continue;
    }

    vector<Cube> cube(n);
    for (int i = 0; i < n; i++) {
      cin >> cube[i].x >> cube[i].y >> cube[i].z;
    }
    Graph g(n);
    for (int i = 0; i < n - 1; i++) {
      for (int j = i + 1; j < n; j++) {
        const ll cost = connected(cube[i], cube[j]);
        if (cost > 0) {
          g.addEdge(i, j, cost);
        }
      }
    }
    vector<bool> visited(n, false);
    vector<Chain> chains;
    for (int i = 0; i < n; i++) {
      if (not visited[i]) {
        vector<bool> visited2(n, false);
        const int dep = getEdge(g, i, visited2);
        Chain chain;
        chain.ring = false;
        dfs(g, dep, dep, visited, chain);
        chains.push_back(chain);
      }
    }
    ll mini = INF;
    for (const auto chain : chains) {
      const bool ring = chain.ring;
      const vector<ll> &costs = chain.edges;
      if (costs.size() < k) {
        continue;
      }
      if (ring) {
        ll sum = 0;
        for (int i = 0; i < k - 1; i++) {
          sum += costs[i];
        }
        if (costs.size() == k) {
          sum += costs[k - 1];
          mini = min(mini, whole - sum);
        } else {
          mini = min(mini, whole - sum);
          int head = 0;
          int tail = k - 1;
          for (int j = 0; j < costs.size();
               j++, head++, tail++, tail = tail % costs.size()) {
            sum += (costs[tail] - costs[head]);
            mini = min(mini, whole - sum);
          }
        }
      } else {
        ll sum = 0;
        for (int i = 0; i < k - 1; i++) {
          sum += costs[i];
        }
        mini = min(mini, whole - sum);
        int head = 0;
        int tail = k - 1;
        for (; tail < costs.size(); head++, tail++) {
          sum += (costs[tail] - costs[head]);
          mini = min(mini, whole - sum);
        }
      }
    }
    cout << ((mini == INF) ? -1 : mini) << endl;
  }
  return 0;
}