Submission #97289

#	Time	Username	Problem	Language	Result	Execution time	Memory
97289		Kemichar	Transport (COCI19_transport)	C++14	130 / 130	305 ms	16784 KiB

Main

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

#define MAXN 100000
#define x first
#define y second

// slightly faster integer input
inline void in(int &N)
{
  N = 0;
  register char ch = 0;
  do {
    ch = getc(stdin);
  } while (ch < '0' || ch > '9');
  while (ch >= '0' && ch <= '9') {
    N = (N << 1) + (N << 3) + ch - '0', ch = getc(stdin);
  }
}

/// input
int iN;
int aiBaseFuel[MAXN];
vector<pair<int, int>> aapiiEdges[MAXN];
/// centroid tree
int aiCentroidLevel[MAXN];
/// temporaries
int aiSubtreeSize[MAXN];
int ctFreeAll;
int actFreeEach[MAXN];
vector<int64_t> allEndsAll;
vector<int64_t> aallEndsEach[MAXN];
pair<int, int> aiiNodeStack[MAXN];

struct State {
  int iCurr;
  int iParent;
  int64_t llBalance;
  int64_t llExtra;
  State(void) {}
  State(int _iCurr, int _iParent, int64_t _llBalance, int64_t _llExtra) :
    iCurr(_iCurr), iParent(_iParent), llBalance(_llBalance), llExtra(_llExtra) {}
};
State aiiStateQueue[MAXN];

// BFS
void CountSubtreeSizes(int iStart)
{
  int iLast = 0;
  aiiNodeStack[0] = make_pair(iStart, iStart);
  for (int i = 0; i <= iLast; i++) {
    for (pair<int, int> &piiEdge : aapiiEdges[aiiNodeStack[i].x]) {
      if (piiEdge.x == aiiNodeStack[i].y || aiCentroidLevel[piiEdge.x] > 0) {
        continue;
      }
      aiiNodeStack[++iLast] = make_pair(piiEdge.x, aiiNodeStack[i].x);
    }
  }
  for (int i = iLast; i >= 0; i--) {
    aiSubtreeSize[aiiNodeStack[i].x] = 1;
    for (pair<int, int> &piiEdge : aapiiEdges[aiiNodeStack[i].x]) {
      if (piiEdge.x == aiiNodeStack[i].y || aiCentroidLevel[piiEdge.x] > 0) {
        continue;
      }
      aiSubtreeSize[aiiNodeStack[i].x] += aiSubtreeSize[piiEdge.x];
    }
  }
}

// finds the centroid for the current subtree and recurses into centroid's children
void FindCentroids_rec(int iStart, int iLevel)
{
  CountSubtreeSizes(iStart);

  const int ctMaxAllowed = aiSubtreeSize[iStart] / 2;
  int iCurr = iStart;
  int iParent = iStart;
  while (true) {
    bool bCentroid = true;
    for (pair<int, int> &piiEdge : aapiiEdges[iCurr]) {
      if (piiEdge.x == iParent || aiCentroidLevel[piiEdge.x] > 0) {
        continue;
      }

      if (aiSubtreeSize[piiEdge.x] > ctMaxAllowed) {
        iParent = iCurr;
        iCurr = piiEdge.x;
        bCentroid = false;
        break;
      }
    }
    if (bCentroid) {
      break;
    }
  }

  aiCentroidLevel[iCurr] = iLevel;
  for (pair<int, int> &piiEdge : aapiiEdges[iCurr]) {
    if (aiCentroidLevel[piiEdge.x] > 0) {
      continue;
    }
    FindCentroids_rec(piiEdge.x, iLevel + 1);
  }
}

// BFS - collects all path costs starting at iStart and ending somewhere in its subtree
void AddPathCosts(int iStart, int iStartParent, int64_t llStartBalance, int64_t llStartExtra, int iSubtree, int iMinCentroidLevel)
{
  aiiStateQueue[0] = State(iStart, iStartParent, llStartBalance, llStartExtra);
  int iQueueLast = 0;
  int iQueueCurr = 0;
  for (int i = 0, j = 0; i <= j; i++) {
    State &sCurr = aiiStateQueue[i];

    if (sCurr.llExtra < 0) {
      sCurr.llBalance += sCurr.llExtra;
      sCurr.llExtra = 0;
    }
    if (sCurr.llBalance >= 0) {
      ctFreeAll++;
      actFreeEach[iSubtree]++;
    } else {
      allEndsAll.push_back(sCurr.llBalance);
      aallEndsEach[iSubtree].push_back(sCurr.llBalance);
    }

    for (pair<int, int> &piiEdge : aapiiEdges[sCurr.iCurr]) {
      if (piiEdge.x == sCurr.iParent || aiCentroidLevel[piiEdge.x] < iMinCentroidLevel) {
        continue;
      }
      aiiStateQueue[++j] = State(piiEdge.x, sCurr.iCurr, sCurr.llBalance, sCurr.llExtra + aiBaseFuel[sCurr.iCurr] - piiEdge.y);
    }
  }
}

// BFS - collects all paths starting in some iStart's subtree and ending in a different iStart's subtree
int64_t CountPaths(int iStart, int iStartParent, int64_t llStartBalance, int64_t llStartExtra, int iSubtree, int iMinCentroidLevel)
{
  int64_t llPaths = 0;

  aiiStateQueue[0] = State(iStart, iStartParent, llStartBalance, llStartExtra);
  int iQueueLast = 0;
  int iQueueCurr = 0;
  for (int i = 0, j = 0; i <= j; i++) {
    State &sCurr = aiiStateQueue[i];

    sCurr.llBalance += aiBaseFuel[sCurr.iCurr];
    if (sCurr.llBalance > 0) {
      sCurr.llExtra += sCurr.llBalance;
      sCurr.llBalance = 0;
    }
    if (sCurr.llBalance == 0) {
      llPaths += allEndsAll.end() - lower_bound(allEndsAll.begin(), allEndsAll.end(), -sCurr.llExtra);
      llPaths -= aallEndsEach[iSubtree].end() - lower_bound(aallEndsEach[iSubtree].begin(), aallEndsEach[iSubtree].end(), -sCurr.llExtra);
      llPaths += ctFreeAll - actFreeEach[iSubtree];
    }

    for (pair<int, int> &piiEdge : aapiiEdges[sCurr.iCurr]) {
      if (piiEdge.x == sCurr.iParent || aiCentroidLevel[piiEdge.x] < iMinCentroidLevel) {
        continue;
      }
      aiiStateQueue[++j] = State(piiEdge.x, sCurr.iCurr, sCurr.llBalance - piiEdge.y, sCurr.llExtra);
    }
  }

  return llPaths;
}

int64_t CountPathsThroughCentroid(int iCentroid)
{
  ctFreeAll = 1;
  allEndsAll.clear();
  int ctSubtrees = 0;
  for (int i = 0; i < aapiiEdges[iCentroid].size(); i++) {
    const pair<int, int> &piiEdge = aapiiEdges[iCentroid][i];
    if (aiCentroidLevel[piiEdge.x] < aiCentroidLevel[iCentroid] + 1) {
      continue;
    }
    actFreeEach[ctSubtrees] = 0;
    aallEndsEach[ctSubtrees].clear();
    AddPathCosts(piiEdge.x, iCentroid, 0, aiBaseFuel[iCentroid] - piiEdge.y, ctSubtrees++, aiCentroidLevel[iCentroid] + 1);
  }
  sort(allEndsAll.begin(), allEndsAll.end());
  for (int i = 0; i < ctSubtrees; i++) {
    sort(aallEndsEach[i].begin(), aallEndsEach[i].end());
  }

  int64_t llPaths = 0;
  for (int i = 0, j = 0; i < aapiiEdges[iCentroid].size(); i++) {
    const pair<int, int> &piiEdge = aapiiEdges[iCentroid][i];
    if (aiCentroidLevel[piiEdge.x] < aiCentroidLevel[iCentroid] + 1) {
      continue;
    }

    llPaths += CountPaths(piiEdge.x, iCentroid, -piiEdge.y, 0, j++, aiCentroidLevel[iCentroid] + 1);
  }
  llPaths += allEndsAll.end() - lower_bound(allEndsAll.begin(), allEndsAll.end(), 0);
  llPaths += ctFreeAll - 1;

  return llPaths;
}

int main(void)
{
  in(iN);
  for (int i = 0; i < iN; i++) {
    in(aiBaseFuel[i]);
  }
  for (int i = 1; i < iN; i++) {
    int iA, iB, iC;
    in(iA); in(iB); in(iC);
    aapiiEdges[iA - 1].push_back(make_pair(iB - 1, iC));
    aapiiEdges[iB - 1].push_back(make_pair(iA - 1, iC));
  }

  FindCentroids_rec(0, 1);

  int64_t llPaths = 0;
  for (int i = 0; i < iN; i++) {
    llPaths += CountPathsThroughCentroid(i);
  }
  printf("%lld\n", llPaths);

  return 0;
}

Compilation message (stderr)

transport.cpp: In function 'void AddPathCosts(int, int, int64_t, int64_t, int, int)':
transport.cpp:113:7: warning: unused variable 'iQueueLast' [-Wunused-variable]
   int iQueueLast = 0;
       ^~~~~~~~~~
transport.cpp:114:7: warning: unused variable 'iQueueCurr' [-Wunused-variable]
   int iQueueCurr = 0;
       ^~~~~~~~~~
transport.cpp: In function 'int64_t CountPaths(int, int, int64_t, int64_t, int, int)':
transport.cpp:145:7: warning: unused variable 'iQueueLast' [-Wunused-variable]
   int iQueueLast = 0;
       ^~~~~~~~~~
transport.cpp:146:7: warning: unused variable 'iQueueCurr' [-Wunused-variable]
   int iQueueCurr = 0;
       ^~~~~~~~~~
transport.cpp: In function 'int64_t CountPathsThroughCentroid(int)':
transport.cpp:177:21: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for (int i = 0; i < aapiiEdges[iCentroid].size(); i++) {
                   ~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
transport.cpp:192:28: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for (int i = 0, j = 0; i < aapiiEdges[iCentroid].size(); i++) {
                          ~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
transport.cpp: In function 'int main()':
transport.cpp:225:27: warning: format '%lld' expects argument of type 'long long int', but argument 2 has type 'int64_t {aka long int}' [-Wformat=]
   printf("%lld\n", llPaths);
                           ^

• Subtask 1: 13 / 13
• Subtask 2: 13 / 13
• Subtask 3: 13 / 13
• Subtask 4: 13 / 13
• Subtask 5: 13 / 13
• Subtask 6: 13 / 13
• Subtask 7: 13 / 13
• Subtask 8: 13 / 13
• Subtask 9: 13 / 13
• Subtask 10: 13 / 13