This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
#include <chrono>
#pragma comment(linker, "/stack:200000000")
#define f first
#define s second
#define pb push_back
using namespace std;
typedef long long ll;
typedef long double ld;
namespace RANDOM
{
ll Time()
{
return chrono::system_clock().now().time_since_epoch().count();
}
mt19937 rnd(Time());
}
using namespace RANDOM;
const int N = 1e5 + 1;
struct tree;
struct tree1;
struct tree1{
int l, r;
tree1 *L, *R;
int sm;
bool used = 1;
void clear()
{
if (used)
{
used = 0;
if (L != NULL) L -> used = 1;
if (R != NULL) R -> used = 1;
sm = 0;
}
}
tree1(int _l, int _r)
{
l = _l, r = _r;
L = NULL;
R = NULL;
sm = 0;
used = 0;
}
void push()
{
if (l == r) return;
int mdl = (l + r) >> 1;
if (L == NULL) L = new tree1(l, mdl);
if (R == NULL) R = new tree1(mdl + 1, r);
}
void insert(int x)
{
push();
clear();
if (l > x || r < x) return;
sm++;
if (l == r) return;
L -> insert(x);
R -> insert(x);
}
int sum(int x)
{
push();
clear();
if (l > x) return 0;
if (r <= x) return sm;
return L -> sum(x) + R -> sum(x);
}
};
struct tree{
int l, r;
tree *L = NULL, *R = NULL;
tree1 *sm = NULL;
bool used = 0;
void clear()
{
if (used)
{
used = 0;
if (L != NULL) L -> used = 1;
if (R != NULL) R -> used = 1;
if (sm != NULL) sm -> used = 1;
}
}
void push()
{
if (sm == NULL) sm = new tree1(0, N);
if (l == r) return;
int mdl = (l + r) >> 1;
if (L == NULL) L = new tree(l, mdl);
if (R == NULL) R = new tree(mdl + 1, r);
}
tree(int _l, int _r)
{
l = _l, r = _r;
push();
}
void insert(int x, int y)
{
clear();
if (l > x || r < x) return;
sm -> insert(y);
if (l == r) return;
L -> insert(x, y);
R -> insert(x, y);
}
int sum(int x, int y)
{
clear();
if (l > x) return 0;
if (r <= x) return sm -> sum(y);
return L -> sum(x, y) + R -> sum(x, y);
}
};
tree *root;
vector <int> a[N];
vector <int> b[N];
int sz[N];
bool mrk[N];
int kol, k;
ll ans = 0;
int Kol(int x, int y)
{
if (mrk[x]) return 0;
sz[x] = 1;
for (int i = 0; i < a[x].size(); i++)
if (a[x][i] != y) sz[x] += Kol(a[x][i], x);
return sz[x];
}
int Find_Centroid(int x, int y)
{
for (int i = 0; i < a[x].size(); i++)
if (a[x][i] != y && !mrk[a[x][i]] && sz[a[x][i]] > kol / 2) return Find_Centroid(a[x][i], x);
return x;
}
void Rec(int x, int y, int sz, int mx, vector <pair<int, int> > *push)
{
if (mrk[x]) return;
push -> pb({sz, mx});
for (int i = 0; i < a[x].size(); i++)
if (a[x][i] != y) Rec(a[x][i], x, sz + 1, max(mx, b[x][i]), push);
}
void Calc_Ans(int x)
{
vector <pair<int, int> > val[a[x].size()];
int m = a[x].size();
for (int i = 0; i < m; i++) Rec(a[x][i], x, 1, b[x][i], &val[i]);
root -> used = 1;
for (int i = 0; i < m; i++)
{
for (int j = 0; j < val[i].size(); j++) ans += root -> sum(val[i][j].s - val[i][j].f - k, val[i][j].s - 1);
for (int j = 0; j < val[i].size(); j++) root -> insert(val[i][j].f, val[i][j].s);
}
root -> used = 1;
for (int i = m - 1; i >= 0; i--)
{
for (int j = 0; j < val[i].size(); j++) ans += root -> sum(val[i][j].s - val[i][j].f - k, val[i][j].s);
for (int j = 0; j < val[i].size(); j++) root -> insert(val[i][j].f, val[i][j].s);
}
for (int i = 0; i < m; i++)
for (int j = 0; j < val[i].size(); j++) ans += ((val[i][j].s - val[i][j].f) >= k);
}
void Centroid(int x)
{
if (mrk[x]) return;
kol = Kol(x, x);
x = Find_Centroid(x, x);
mrk[x] = 1;
Calc_Ans(x);
for (int i = 0; i < a[x].size(); i++) Centroid(a[x][i]);
}
int main()
{
root = new tree(0, N);
int n;
cin >> n >> k;
for (int i = 0; i < n - 1; i++)
{
int x, y, z;
cin >> x >> y >> z;
a[--x].pb(--y);
a[y].pb(x);
b[x].pb(z);
b[y].pb(z);
}
Centroid(0);
cout << ans * 2;
}
Compilation message (stderr)
Main.cpp:3: warning: ignoring #pragma comment [-Wunknown-pragmas]
3 | #pragma comment(linker, "/stack:200000000")
|
Main.cpp: In function 'int Kol(int, int)':
Main.cpp:148:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
148 | for (int i = 0; i < a[x].size(); i++)
| ~~^~~~~~~~~~~~~
Main.cpp: In function 'int Find_Centroid(int, int)':
Main.cpp:155:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
155 | for (int i = 0; i < a[x].size(); i++)
| ~~^~~~~~~~~~~~~
Main.cpp: In function 'void Rec(int, int, int, int, std::vector<std::pair<int, int> >*)':
Main.cpp:164:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
164 | for (int i = 0; i < a[x].size(); i++)
| ~~^~~~~~~~~~~~~
Main.cpp: In function 'void Calc_Ans(int)':
Main.cpp:176:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
176 | for (int j = 0; j < val[i].size(); j++) ans += root -> sum(val[i][j].s - val[i][j].f - k, val[i][j].s - 1);
| ~~^~~~~~~~~~~~~~~
Main.cpp:177:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
177 | for (int j = 0; j < val[i].size(); j++) root -> insert(val[i][j].f, val[i][j].s);
| ~~^~~~~~~~~~~~~~~
Main.cpp:182:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
182 | for (int j = 0; j < val[i].size(); j++) ans += root -> sum(val[i][j].s - val[i][j].f - k, val[i][j].s);
| ~~^~~~~~~~~~~~~~~
Main.cpp:183:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
183 | for (int j = 0; j < val[i].size(); j++) root -> insert(val[i][j].f, val[i][j].s);
| ~~^~~~~~~~~~~~~~~
Main.cpp:186:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
186 | for (int j = 0; j < val[i].size(); j++) ans += ((val[i][j].s - val[i][j].f) >= k);
| ~~^~~~~~~~~~~~~~~
Main.cpp: In function 'void Centroid(int)':
Main.cpp:196:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
196 | for (int i = 0; i < a[x].size(); i++) Centroid(a[x][i]);
| ~~^~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |