#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
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);
^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
11 ms |
5504 KB |
Output is correct |
| 2 |
Correct |
12 ms |
5508 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
12 ms |
5504 KB |
Output is correct |
| 2 |
Correct |
12 ms |
5732 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
59 ms |
9828 KB |
Output is correct |
| 2 |
Correct |
53 ms |
9452 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
84 ms |
11704 KB |
Output is correct |
| 2 |
Correct |
106 ms |
12816 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
121 ms |
14072 KB |
Output is correct |
| 2 |
Correct |
156 ms |
15976 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
62 ms |
7768 KB |
Output is correct |
| 2 |
Correct |
32 ms |
7844 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
59 ms |
10100 KB |
Output is correct |
| 2 |
Correct |
115 ms |
10128 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
155 ms |
11560 KB |
Output is correct |
| 2 |
Correct |
161 ms |
12236 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
262 ms |
13296 KB |
Output is correct |
| 2 |
Correct |
238 ms |
14060 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
305 ms |
16784 KB |
Output is correct |
| 2 |
Correct |
273 ms |
15992 KB |
Output is correct |
