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

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <map>
#include <vector>
#define INF 1000001LL
#define INF2 100000001000000LL
using namespace std;

typedef long long ll;

struct cord {
  int d, k;
  vector<ll> c;
  cord(void) {
    d = k = 0;
    c.clear();
  }
  cord(vector<ll> v) {
    d = v.size();
    k = 0;
    c = v;
  }
  bool operator<(const cord a) const { return c[k] < a.c[k]; }
  bool operator==(const cord a) const { return c == a.c; }
};

ll dis(cord a, cord b) {
  if (a.d != b.d)
    return INF;
  ll res = 0;
  for (int i = 0; i < a.d; i++)
    res += (a.c[i] - b.c[i]) * (a.c[i] - b.c[i]);
  return res;
}

struct KDTree {
  int n, d, root;
  vector<cord> v;

  KDTree(void) {
    n = d = root = 0;
    v.clear();
  }

  KDTree(vector<cord> x) {
    n = x.size();
    if (!n)
      return;
    d = x[0].c.size();
    if (!d)
      return;
    v = x;
    compose(0, n, 0);
  }

  void compose(int l, int r, int depth) {
    if (r - l <= 1)
      return;
    for (int i = l; i < r; i++)
      v[i].k = depth;
    sort(v.begin() + l, v.begin() + r);

    compose(l, (l + r) / 2, (depth + 1) % d);
    compose((l + r) / 2 + 1, r, (depth + 1) % d);
  }

  inline void search(int l, int r, vector<ll> &rng, cord p, int depth,
                     ll &dist) {
    if (r - l < 1)
      return;

    int mid = (l + r) / 2;
    dist = min(dist, dis(p, v[mid]));
    for (int i = 0; i < d; i++) {
      if (rng[2 * i] > p.c[i] &&
          (rng[2 * i] - p.c[i]) * (rng[2 * i] - p.c[i]) >= dist)
        return;
      else if (rng[2 * i + 1] < p.c[i] &&
               (rng[2 * i + 1] - p.c[i]) * (rng[2 * i + 1] - p.c[i]) >= dist)
        return;
    }

    ll tmp, nxt[4], st[] = {0, 1}, d1, d2;
    nxt[0] = mid + 1;
    nxt[1] = r;
    nxt[2] = l;
    nxt[3] = mid;

    if (p.c[depth] < v[mid].c[depth]) {
      swap(nxt[0], nxt[2]);
      swap(nxt[1], nxt[3]);
      swap(st[0], st[1]);
    }

    for (int i = 0; i < 2; i++) {
      tmp = rng[2 * depth + st[i]];
      rng[2 * depth + st[i]] = v[mid].c[depth];
      search(nxt[2 * i], nxt[2 * i + 1], rng, p, (depth + 1) % d, dist);
      rng[2 * depth + st[i]] = tmp;
    }
  }
};

int n, q, m;
vector<cord> v;
cord p, r;
vector<ll> tmp;

int main() {
  vector<ll> range;
  for (int i = 0; i < 3; i++) {
    range.push_back(-INF);
    range.push_back(INF);
  }

  while (scanf("%d%d", &n, &q) != EOF, n || q) {
    v.resize(n);
    tmp.resize(3);

    for (int i = 0; i < n; i++) {
      for (int j = 0; j < 3; j++)
        scanf("%lld", &tmp[j]);
      v[i] = cord(tmp);
    }

    KDTree kd(v);

    for (int i = 0; i < q; i++) {
      scanf("%d", &m);
      double ans = 0;
      ll now, prv;

      for (int j = 0; j < 3; j++)
        scanf("%lld", &tmp[j]);
      p = cord(tmp);
      prv = INF2;
      kd.search(0, n, range, p, 0, prv);
      for (int j = 1; j < m; j++) {
        for (int k = 0; k < 3; k++)
          scanf("%lld", &tmp[k]);

        r = cord(tmp);
        now = INF2;
        kd.search(0, n, range, r, 0, now);
        ans += min(sqrt(prv) + sqrt(now), sqrt(dis(p, r)));
        prv = now;
        p = r;
      }
      printf("%.9lf\\n", ans);
    }
  }
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <ctime>
#include <iostream>
#include <map>
#include <vector>
#define INF 1000001LL
#define INF2 100000001000000LL
using namespace std;

typedef long long ll;
int cnt;
struct cord {
  int d, k;
  vector<ll> c;
  cord(void) {
    d = k = 0;
    c.clear();
  }
  cord(vector<ll> v) {
    d = v.size();
    k = 0;
    c = v;
  }
  bool operator<(const cord a) const { return c[k] < a.c[k]; }
  bool operator==(const cord a) const { return c == a.c; }
};

ll dis(cord a, cord b) {
  if (a.d != b.d)
    return INF;
  ll res = 0;
  for (int i = 0; i < a.d; i++)
    res += (a.c[i] - b.c[i]) * (a.c[i] - b.c[i]);
  return res;
}

struct KDTree {
  int n, d, root;
  vector<cord> v;

  KDTree(void) {
    n = d = root = 0;
    v.clear();
  }

  KDTree(vector<cord> x) {
    n = x.size();
    if (!n)
      return;
    d = x[0].c.size();
    if (!d)
      return;
    v = x;
    compose(0, n, 0);
  }

  void compose(int l, int r, int depth) {
    if (r - l <= 1)
      return;
    for (int i = l; i < r; i++)
      v[i].k = depth;
    sort(v.begin() + l, v.begin() + r);

    compose(l, (l + r) / 2, (depth + 1) % d);
    compose((l + r) / 2 + 1, r, (depth + 1) % d);
  }

  void search(int l, int r, vector<ll> &rng, cord p, int depth, ll &dist) {
    if (r - l < 1)
      return;
    cnt++;
    for (int i = 0; i < d; i++) {
      if (rng[2 * i] > p.c[i] &&
          (rng[2 * i] - p.c[i]) * (rng[2 * i] - p.c[i]) > dist)
        return;
      else if (rng[2 * i + 1] < p.c[i] &&
               (rng[2 * i + 1] - p.c[i]) * (rng[2 * i + 1] - p.c[i]) > dist)
        return;
    }

    int mid = (l + r) / 2;
    dist = min(dist, dis(p, v[mid]));

    ll tmp, nxt[4], st[] = {0, 1}, d1, d2;
    nxt[0] = mid + 1;
    nxt[1] = r;
    nxt[2] = l;
    nxt[3] = mid;
    if (nxt[1] - nxt[0] < 1)
      d1 = INF;
    else
      d1 = dis(p, v[(nxt[0] + nxt[1]) / 2]);
    if (nxt[3] - nxt[2] < 1)
      d2 = INF;
    else
      d2 = dis(p, v[(nxt[2] + nxt[3]) / 2]);

    if (d1 > d2) {
      swap(nxt[0], nxt[2]);
      swap(nxt[1], nxt[3]);
      swap(st[0], st[1]);
    }

    for (int i = 0; i < 2; i++) {
      tmp = rng[2 * depth + st[i]];
      rng[2 * depth + st[i]] = v[mid].c[depth];
      search(nxt[2 * i], nxt[2 * i + 1], rng, p, (depth + 1) % d, dist);
      rng[2 * depth + st[i]] = tmp;
    }
  }
};

int n, q, m;
vector<cord> v;
cord p, r;
vector<ll> tmp;

int main() {
  vector<ll> range;
  for (int i = 0; i < 3; i++) {
    range.push_back(-INF);
    range.push_back(INF);
  }

  while (scanf("%d%d", &n, &q) != EOF, n || q) {
    v.resize(n);
    tmp.resize(3);

    for (int i = 0; i < n; i++) {
      for (int j = 0; j < 3; j++)
        scanf("%lld", &tmp[j]);
      v[i] = cord(tmp);
    }

    KDTree kd(v);

    for (int i = 0; i < q; i++) {
      scanf("%d", &m);
      double ans = 0;
      ll now, prv;

      for (int j = 0; j < 3; j++)
        scanf("%lld", &tmp[j]);
      cnt = 0;
      p = cord(tmp);
      prv = INF2;
      kd.search(0, n, range, p, 0, prv);
      // cout << cnt << endl;
      for (int j = 1; j < m; j++) {
        for (int k = 0; k < 3; k++)
          scanf("%lld", &tmp[k]);

        cnt = 0;
        r = cord(tmp);
        now = INF2;
        kd.search(0, n, range, r, 0, now);
        // cout << cnt << endl;
        ans += min(sqrt(prv) + sqrt(now), sqrt(dis(p, r)));
        prv = now;
        p = r;
      }
      printf("%.9lf\\n", ans);
    }
  }
}