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 <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key
typedef long long ll;
typedef pair<ll,ll> ii;
typedef vector<ll> vi;
typedef long double ld;
typedef tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update> pbds;
typedef set<int>::iterator sit;
typedef map<int,int>::iterator mit;
typedef vector<int>::iterator vit;
struct edge
{
int v; ll cost, val; //cost means cost of change
int lab;
};
vector<vector<edge> > adj;
vi par;
vi ans;
vi cost;
vi c;
ll res;
vi record;
vi record2;
vector<bool> isleaf;
edge edg(int v, ll cost, ll val, int lab)
{
edge tmp; tmp.v = v; tmp.cost = cost; tmp.val = val; tmp.lab = lab;
return tmp;
}
void dfs(int u, int p)
{
par[u] = p;
if(p==-1) c[u] = -1;
for(int i = 0; i < adj[u].size(); i++)
{
if(adj[u][i].v == p) continue;
c[adj[u][i].v] = adj[u][i].val;
cost[adj[u][i].v] = adj[u][i].cost;
dfs(adj[u][i].v, u);
}
}
void popmax(map<ll,ll>* pq, ll k) //pops the k largest slopes
{
map<ll,ll>::iterator it = pq->end();
while(k && pq->size()>0)
{
it = pq->end();
it--;
ll occ = it->se;
if(occ>k)
{
pq->at(it->fi) -= k;
k = 0;
}
else
{
pq->erase(it->fi);
k -= occ;
}
}
}
void popmin(map<ll,ll>* pq, ll k) //pops the k smallest slopes
{
map<ll,ll>::iterator it = pq->begin();
while(k && pq->size()>0)
{
it = pq->begin();
ll occ = it->se;
if(occ>k)
{
pq->at(it->fi) -= k;
k = 0;
}
else
{
pq->erase(it->fi);
k -= occ;
}
}
}
ii getmax(map<ll,ll>* pq)
{
map<ll,ll>::iterator it = pq->end();
it--;
return mp(it->fi,it->se);
}
ii getmin(map<ll,ll>* pq)
{
map<ll,ll>::iterator it = pq->begin();
return mp(it->fi,it->se);
}
struct data
{
ll a,b;
map<ll,ll> *pq;
ll shift;
data operator+(data r){
data s;
s.a=a+r.a;
s.b=b+r.b;
s.shift = 0;
if(pq->size()>r.pq->size())
{
s.pq=pq;
s.shift = shift;
ll change = r.shift-s.shift;
while(r.pq->size()!=0)
{
ii tmp = getmax(r.pq);
pq->insert(mp(tmp.fi+change,0));
pq->at(tmp.fi+change) += tmp.se;
r.pq->erase(tmp.fi);
}
}
else
{
s.pq=r.pq;
s.shift = r.shift;
ll change = shift-s.shift;
while(pq->size()!=0)
{
ii tmp = getmax(pq);
r.pq->insert(mp(tmp.fi+change,0));
r.pq->at(tmp.fi+change) += tmp.se;
pq->erase(tmp.fi);
}
}
return s;
}
};
vector<data> d;
vector<bool> isgood;
void procnode(data &s, ll c, ll d)
{
ll ext = max(0LL, s.a - d);
ll k = ext;
map<ll,ll>::iterator it = s.pq->end();
while(k && s.pq->size()>0)
{
it = s.pq->end();
it--;
ll occ = it->se;
if(occ>k)
{
s.a -= k;
s.b += (it->fi+s.shift)*k;
s.pq->at(it->fi) -= k;
k = 0;
}
else
{
s.a -= occ;
s.b += (it->fi+s.shift)*occ;
s.pq->erase(it->fi);
k -= occ;
}
}
popmin(s.pq, ext);
s.shift += c;
s.b -= c*s.a;
}
void calc(int u, int p)
{
bool visit = false;
d[u].a=0;
d[u].b=0;
d[u].shift=0;
d[u].pq = new map<ll,ll>;
for(int i = 0; i < adj[u].size(); i++)
{
if(adj[u][i].v == p) continue;
visit = true;
calc(adj[u][i].v, u);
}
if(p==-1) return ;
if(visit)
{
ll cc = c[u]; ll dd = cost[u];
if(d[u].a<dd) isgood[u] = true;
procnode(d[u], cc, dd);
record[u] = (getmin(d[u].pq).fi+d[u].shift);
record2[u] = (getmax(d[u].pq).fi+d[u].shift);
d[p] = d[p]+d[u];
}
else
{
isleaf[u] = true;
d[u].a = cost[u];
d[u].b = -c[u]*cost[u];
d[u].shift = 0;
d[u].pq->insert(mp(c[u],0));
d[u].pq->at(c[u]) += 2*cost[u];
record[u] = (getmin(d[u].pq).fi+d[u].shift);
record2[u] = (getmax(d[u].pq).fi+d[u].shift);
d[p] = d[p]+d[u];
}
}
void construct(int u, int p, ll s)
{
for(int i = 0; i < adj[u].size(); i++)
{
if(adj[u][i].v == p) continue;
int v = adj[u][i].v;
int lab = adj[u][i].lab;
ll cur = 0; //stores the up edge length
ll l = record[v]; ll r = record2[v];
if(!isgood[v])
{
if(l<=s&&s<=r)
{
cur = c[v];
}
else if(s<l)
{
cur = c[v] - (l - s);
}
else
{
cur = c[v] + (s - r);
}
}
else
{
cur = c[v];
}
ans[lab] = cur;
construct(v, u, s-cur);
}
}
ll change = 0;
bool check(int u, int p, ll s)
{
bool answ = 0;
for(int i = 0; i < adj[u].size(); i++)
{
if(adj[u][i].v == p) continue;
int v = adj[u][i].v;
int lab = adj[u][i].lab;
change += abs(c[v] - ans[lab])*cost[v];
answ |= check(v,u,s-ans[lab]);
}
if(isleaf[u])
{
if(s==0) return true;
else return false;
}
return answ;
}
int main()
{
ios_base::sync_with_stdio(0); cin.tie(0);
int t; cin >> t;
while(t--)
{
{
adj.clear();
par.clear();
ans.clear();
d.clear();
c.clear();
record.clear();
cost.clear();
isgood.clear();
isleaf.clear();
record2.clear();
}
int n; cin >> n;
vector<ii> vec;
ll tot = 0;
adj.resize(n); par.resize(n); ans.resize(n); d.resize(n); c.resize(n);
record.resize(n); isleaf.assign(n,0); record2.resize(n); cost.resize(n); isgood.assign(n,0);
for(int i = 0; i < n - 1; i++)
{
int u, v; cin>>u>>v;
u--; v--;
ll cc, dd; cin>>cc;
cin>>dd;
adj[u].pb(edg(v,dd,cc,i));
adj[v].pb(edg(u,dd,cc,i));
vec.pb(mp(cc,dd));
tot+=dd;
}
dfs(0,-1);
calc(0,-1);
while(d[0].a>0 && d[0].pq->size()>0)
{
map<ll,ll>::iterator it = d[0].pq->end();
it--;
ll occ = it->se;
if(occ>d[0].a)
{
d[0].b += (it->fi+d[0].shift)*d[0].a;
d[0].pq->at(it->fi) -= d[0].a;
d[0].a = 0;
}
else
{
d[0].b += (it->fi+d[0].shift)*occ;
d[0].pq->erase(it->fi);
d[0].a -= occ;
}
}
res = d[0].b;
record[0] = getmax(d[0].pq).fi+d[0].shift;
construct(0,-1,record[0]);
change=0;
bool correct = check(0,-1,record[0]);
assert(correct&&(change==res));
cout<<res<<'\n';
for(int i = 0; i < n - 1; i++)
{
cout << ans[i] << '\n';
}
}
}
Compilation message (stderr)
fireworks.cpp: In function 'void dfs(int, int)':
fireworks.cpp:50:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<edge>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
50 | for(int i = 0; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
fireworks.cpp: At global scope:
fireworks.cpp:153:12: error: template argument 1 is invalid
153 | vector<data> d;
| ^
fireworks.cpp:153:12: error: template argument 2 is invalid
fireworks.cpp:156:6: error: variable or field 'procnode' declared void
156 | void procnode(data &s, ll c, ll d)
| ^~~~~~~~
fireworks.cpp:156:15: error: reference to 'data' is ambiguous
156 | void procnode(data &s, ll c, ll d)
| ^~~~
In file included from /usr/include/c++/10/string:54,
from /usr/include/c++/10/bits/locale_classes.h:40,
from /usr/include/c++/10/bits/ios_base.h:41,
from /usr/include/c++/10/ios:42,
from /usr/include/c++/10/istream:38,
from /usr/include/c++/10/sstream:38,
from /usr/include/c++/10/complex:45,
from /usr/include/c++/10/ccomplex:39,
from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:54,
from fireworks.cpp:1:
/usr/include/c++/10/bits/range_access.h:319:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(std::initializer_list<_Tp>)'
319 | data(initializer_list<_Tp> __il) noexcept
| ^~~~
/usr/include/c++/10/bits/range_access.h:310:5: note: 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
310 | data(_Tp (&__array)[_Nm]) noexcept
| ^~~~
/usr/include/c++/10/bits/range_access.h:300:5: note: 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
300 | data(const _Container& __cont) noexcept(noexcept(__cont.data()))
| ^~~~
/usr/include/c++/10/bits/range_access.h:290:5: note: 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
290 | data(_Container& __cont) noexcept(noexcept(__cont.data()))
| ^~~~
fireworks.cpp:113:8: note: 'struct data'
113 | struct data
| ^~~~
fireworks.cpp:156:21: error: 's' was not declared in this scope
156 | void procnode(data &s, ll c, ll d)
| ^
fireworks.cpp:156:27: error: expected primary-expression before 'c'
156 | void procnode(data &s, ll c, ll d)
| ^
fireworks.cpp:156:33: error: expected primary-expression before 'd'
156 | void procnode(data &s, ll c, ll d)
| ^
fireworks.cpp: In function 'void calc(int, int)':
fireworks.cpp:189:3: error: invalid types 'int[int]' for array subscript
189 | d[u].a=0;
| ^
fireworks.cpp:190:3: error: invalid types 'int[int]' for array subscript
190 | d[u].b=0;
| ^
fireworks.cpp:191:3: error: invalid types 'int[int]' for array subscript
191 | d[u].shift=0;
| ^
fireworks.cpp:192:3: error: invalid types 'int[int]' for array subscript
192 | d[u].pq = new map<ll,ll>;
| ^
fireworks.cpp:193:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<edge>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
193 | for(int i = 0; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
fireworks.cpp:203:7: error: invalid types 'int[int]' for array subscript
203 | if(d[u].a<dd) isgood[u] = true;
| ^
fireworks.cpp:204:13: error: invalid types 'int[int]' for array subscript
204 | procnode(d[u], cc, dd);
| ^
fireworks.cpp:204:3: error: 'procnode' was not declared in this scope
204 | procnode(d[u], cc, dd);
| ^~~~~~~~
fireworks.cpp:205:24: error: invalid types 'int[int]' for array subscript
205 | record[u] = (getmin(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:205:36: error: invalid types 'int[int]' for array subscript
205 | record[u] = (getmin(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:206:25: error: invalid types 'int[int]' for array subscript
206 | record2[u] = (getmax(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:206:37: error: invalid types 'int[int]' for array subscript
206 | record2[u] = (getmax(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:207:4: error: invalid types 'int[int]' for array subscript
207 | d[p] = d[p]+d[u];
| ^
fireworks.cpp:207:11: error: invalid types 'int[int]' for array subscript
207 | d[p] = d[p]+d[u];
| ^
fireworks.cpp:207:16: error: invalid types 'int[int]' for array subscript
207 | d[p] = d[p]+d[u];
| ^
fireworks.cpp:212:4: error: invalid types 'int[int]' for array subscript
212 | d[u].a = cost[u];
| ^
fireworks.cpp:213:4: error: invalid types 'int[int]' for array subscript
213 | d[u].b = -c[u]*cost[u];
| ^
fireworks.cpp:214:4: error: invalid types 'int[int]' for array subscript
214 | d[u].shift = 0;
| ^
fireworks.cpp:215:4: error: invalid types 'int[int]' for array subscript
215 | d[u].pq->insert(mp(c[u],0));
| ^
fireworks.cpp:216:4: error: invalid types 'int[int]' for array subscript
216 | d[u].pq->at(c[u]) += 2*cost[u];
| ^
fireworks.cpp:217:24: error: invalid types 'int[int]' for array subscript
217 | record[u] = (getmin(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:217:36: error: invalid types 'int[int]' for array subscript
217 | record[u] = (getmin(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:218:25: error: invalid types 'int[int]' for array subscript
218 | record2[u] = (getmax(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:218:37: error: invalid types 'int[int]' for array subscript
218 | record2[u] = (getmax(d[u].pq).fi+d[u].shift);
| ^
fireworks.cpp:219:4: error: invalid types 'int[int]' for array subscript
219 | d[p] = d[p]+d[u];
| ^
fireworks.cpp:219:11: error: invalid types 'int[int]' for array subscript
219 | d[p] = d[p]+d[u];
| ^
fireworks.cpp:219:16: error: invalid types 'int[int]' for array subscript
219 | d[p] = d[p]+d[u];
| ^
fireworks.cpp: In function 'void construct(int, int, ll)':
fireworks.cpp:225:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<edge>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
225 | for(int i = 0; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
fireworks.cpp: In function 'bool check(int, int, ll)':
fireworks.cpp:260:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<edge>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
260 | for(int i = 0; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
fireworks.cpp: In function 'int main()':
fireworks.cpp:286:6: error: request for member 'clear' in 'd', which is of non-class type 'int'
286 | d.clear();
| ^~~~~
fireworks.cpp:297:50: error: request for member 'resize' in 'd', which is of non-class type 'int'
297 | adj.resize(n); par.resize(n); ans.resize(n); d.resize(n); c.resize(n);
| ^~~~~~
fireworks.cpp:312:10: error: invalid types 'int[int]' for array subscript
312 | while(d[0].a>0 && d[0].pq->size()>0)
| ^
fireworks.cpp:312:22: error: invalid types 'int[int]' for array subscript
312 | while(d[0].a>0 && d[0].pq->size()>0)
| ^
fireworks.cpp:314:31: error: invalid types 'int[int]' for array subscript
314 | map<ll,ll>::iterator it = d[0].pq->end();
| ^
fireworks.cpp:317:12: error: invalid types 'int[int]' for array subscript
317 | if(occ>d[0].a)
| ^
fireworks.cpp:319:6: error: invalid types 'int[int]' for array subscript
319 | d[0].b += (it->fi+d[0].shift)*d[0].a;
| ^
fireworks.cpp:319:24: error: invalid types 'int[int]' for array subscript
319 | d[0].b += (it->fi+d[0].shift)*d[0].a;
| ^
fireworks.cpp:319:36: error: invalid types 'int[int]' for array subscript
319 | d[0].b += (it->fi+d[0].shift)*d[0].a;
| ^
fireworks.cpp:320:6: error: invalid types 'int[int]' for array subscript
320 | d[0].pq->at(it->fi) -= d[0].a;
| ^
fireworks.cpp:320:29: error: invalid types 'int[int]' for array subscript
320 | d[0].pq->at(it->fi) -= d[0].a;
| ^
fireworks.cpp:321:6: error: invalid types 'int[int]' for array subscript
321 | d[0].a = 0;
| ^
fireworks.cpp:325:6: error: invalid types 'int[int]' for array subscript
325 | d[0].b += (it->fi+d[0].shift)*occ;
| ^
fireworks.cpp:325:24: error: invalid types 'int[int]' for array subscript
325 | d[0].b += (it->fi+d[0].shift)*occ;
| ^
fireworks.cpp:326:6: error: invalid types 'int[int]' for array subscript
326 | d[0].pq->erase(it->fi);
| ^
fireworks.cpp:327:6: error: invalid types 'int[int]' for array subscript
327 | d[0].a -= occ;
| ^
fireworks.cpp:330:10: error: invalid types 'int[int]' for array subscript
330 | res = d[0].b;
| ^
fireworks.cpp:331:23: error: invalid types 'int[int]' for array subscript
331 | record[0] = getmax(d[0].pq).fi+d[0].shift;
| ^
fireworks.cpp:331:35: error: invalid types 'int[int]' for array subscript
331 | record[0] = getmax(d[0].pq).fi+d[0].shift;
| ^