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

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define ALL(A) A.begin(), A.end()

using namespace std;

typedef long long ll;
typedef pair<int, int> P;

const int MAX_N = 2 * (int)1e5 + 5;

struct Graph {
public:
  int to;
  int cost;
  int depth;
};

struct XYC {
public:
  int x, y, c;

  XYC() {}

  void setX(int _x) { x = _x; }

  void setY(int _y) { y = _y; }

  void setC(int _c) { c = _c; }

  int getX(void) { return x; }

  int getY(void) { return y; }

  int getC(void) { return c; }

  bool operator<(const XYC &xyc) const {
    return (x != xyc.x ? x < xyc.x : y < xyc.y);
  }
};

vector<Graph> G[MAX_N];

int sumCost[MAX_N]; // ?????±???????????????????????????????¨?
int numNode[MAX_N]; // ?????±?????????????????????????????°
int maxDepth;       // ?????§??±???

bool used[MAX_N];

void dfs(int curr, int depth, int prev) {
  if (used[curr])
    return;
  used[curr] |= true;

  for (int i = 0; i < G[curr].size(); ++i) {
    int next = G[curr][i].to;
    if (!used[next] && next != prev) {
      G[curr][i].depth = depth;
      maxDepth = max(maxDepth, depth);
      dfs(next, depth + 1, curr);
    } // end if
  }   // end for
}

void dfs2(int curr, int depth, int prev) {
  if (used[curr])
    return;
  used[curr] |= true;

  for (int i = 0; i < G[curr].size(); ++i) {
    int next = G[curr][i].to;
    if (!used[next] && next != prev) {
      sumCost[depth] += G[curr][i].cost;
      ++numNode[depth];
      dfs2(next, depth + 1, curr);
    } // edn if
  }   // end for
}

int main() {
  memset(sumCost, 0., sizeof(sumCost));
  memset(numNode, 0, sizeof(numNode));
  maxDepth = 0;

  rep(i, MAX_N) G[i].clear();
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  double p;
  cin >> p;
  int N;
  cin >> N;
  vector<XYC> xyc(N - 1);

  rep(i, N - 1) {
    int x, y, c;
    cin >> x >> y >> c;
    --x, --y;
    xyc[i].setX(x);
    xyc[i].setY(y);
    xyc[i].setC(c);
  } // end rep

  sort(ALL(xyc));

  rep(i, N - 1) {
    G[xyc[i].getX()].push_back((Graph){xyc[i].getY(), xyc[i].getC(), 0});
    G[xyc[i].getY()].push_back((Graph){xyc[i].getX(), xyc[i].getC(), 0});
  } // end rep

  memset(used, false, sizeof(used));
  dfs(0, 1, -1);
  memset(used, false, sizeof(used));
  dfs2(0, 1, -1);

  double orig = 0.;
  double currp = p;
  for (int i = 1; i <= maxDepth; ++i) {
    orig += (double)sumCost[i] * currp;
    currp *= p;
  } // end for

  double res = orig;
  //	cerr << "orig: " << orig << endl;
  currp = p;
  for (int k = 1; numNode[k] != 0; ++k) {
    res += sumCost[k] * currp;
    res += numNode[k] * currp * orig;
    currp *= p;
    //		cerr << "numNode: " << numNode[k] << " sumCost: " << sumCost[k] <<
    //endl;
  } // end for

  printf("%.8lf\\n", res);

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define ALL(A) A.begin(), A.end()

using namespace std;

typedef long long ll;
typedef pair<int, int> P;

const int MAX_N = 2 * (int)1e5 + 5;

struct Graph {
public:
  int to;
  int cost;
  int depth;
};

vector<Graph> G[MAX_N];

int Depth[MAX_N];   // ?????±???????????????????????????????¨?
int numNode[MAX_N]; // ?????±?????????????????????????????°
int maxDepth;       // ?????§??±???
bool used[MAX_N];

void dfs(int curr, int depth) {
  used[curr] = true;
  for (int i = 0; i < G[curr].size(); ++i) {
    G[curr][i].depth = depth;
    maxDepth = max(maxDepth, depth);
    Depth[depth] += G[curr][i].cost;
    ++numNode[depth];
  } // end for

  for (int i = 0; i < G[curr].size(); ++i) {
    int next = G[curr][i].to;
    if (!used[next])
      dfs(next, depth + 1);
  } // end for
}

int sum[MAX_N];

int main() {
  memset(Depth, 0, sizeof(Depth));
  memset(numNode, 0, sizeof(numNode));
  maxDepth = 0;

  rep(i, MAX_N) G[i].clear();
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  double p;
  cin >> p;
  int N;
  cin >> N;
  vector<int> x(N - 1, 0);
  vector<int> y(N - 1, 0);
  vector<int> c(N - 1, 0);
  rep(i, N - 1) {
    cin >> x[i] >> y[i] >> c[i];
    --x[i], --y[i];
    G[x[i]].push_back((Graph){y[i], c[i], 0});
  } // end rep
  numNode[0] = 1;
  memset(used, false, sizeof(used));
  dfs(0, 1);

  memset(sum, 0, sizeof(sum));
  for (int i = 0; i <= maxDepth; ++i) {
    sum[i] += Depth[i];
    for (int k = 1; k <= maxDepth; ++k) {
      sum[i + k] += numNode[i] * Depth[k];
    } // end for
  }   // end for

  double currp = p;
  double res = 0.;
  for (int i = 1; i < MAX_N; ++i) {
    res += (double)sum[i] * currp;
    currp *= p;
  } // end for

  printf("%.8lf\\n", res);

  return 0;
}