-
Notifications
You must be signed in to change notification settings - Fork 44
/
Copy pathmain.cpp
163 lines (132 loc) · 3.72 KB
/
main.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
#include <bits/stdc++.h>
using namespace std; // NOLINT
//
// Represents a node in the tree.
//
struct node {
int weight;
vector<int> edges;
};
//
// Represents an edge from x to y with weight w and strength p.
//
struct edge {
int x;
int y;
int w;
int p;
};
//
// Parses an edge from the istream.
//
istream& operator>>(istream& is, edge& e) {
is >> e.x >> e.y >> e.w >> e.p;
e.x--;
e.y--;
return is;
}
//
// Writes an edge to the ostream.
//
ostream& operator<<(ostream& os, const edge& e) {
return os << e.x + 1 << " " << e.y + 1 << " " << e.w << " " << e.p;
}
//
// Lower the weight of the tree as much as possible while fixing it using
// DFS. Returns false if the tree cannot be fixed and true otherwise.
//
bool lower(vector<node>& nodes, vector<edge>& edges, int i) {
node& n = nodes[i];
n.weight = 0;
if (n.edges.empty()) {
return true;
}
for (int j : n.edges) {
edge& e = edges[j];
node& child = nodes[e.y];
if (!lower(nodes, edges, e.y) || e.p < child.weight) {
return false;
}
// Now that the child this edge is connected to has been lowered, lower
// the weight of the edge as much as possible.
int cut = min(e.w - 1, e.p - child.weight);
e.w -= cut;
e.p -= cut;
n.weight += e.w + child.weight;
}
return true;
}
//
// Raises the weight of the tree up to w_upper weight while keeping it fixed.
// We do this by processing each (edge, child) pair of the current node by
// first increasing the weight of the edge as much as possible and then using
// that to determine an upper bound for the recursive procedur on the child.
//
void raise(vector<node>& nodes,
vector<edge>& edges,
const vector<edge>& original,
int i,
int w_upper) {
assert(w_upper >= 0);
node& n = nodes[i];
if (n.edges.empty() || n.weight == w_upper) {
// No way to increase weight or upper weight bound has been met.
return;
}
for (int j : n.edges) {
edge& e = edges[j];
node& child = nodes[e.y];
// Figure out the most we can increase the weight of the edge to the
// current child by. This is only limited by the original edge weight
// for edges stemming from the root.
int e_w_increase = original[j].w - e.w;
if (i != 0) {
e_w_increase = min(e_w_increase, w_upper - n.weight);
}
// Increase the weight and strength of the edge.
e.w += e_w_increase;
e.p += e_w_increase;
n.weight += e_w_increase;
// Figure out the upper bound on the weight of the child that would keep
// the tree fixed. Once again, the root can have unlimited weight in which
// case w_upper does not apply.
int c_w_upper = e.p;
if (i != 0) {
c_w_upper = min(c_w_upper, w_upper - n.weight + child.weight);
}
int c_w_original = child.weight;
raise(nodes, edges, original, e.y, c_w_upper);
// Update the weight of this node.
n.weight += (child.weight - c_w_original);
}
}
//
// The overall strategy being used for this problem is a 2-phase DFS.
//
// 1) Lower the weight of the tree as much as possible while fixing it. Detect
// if the tree is not fixable and terminate if needed.
//
// 2) Raise the weight of the lowered and fixed tree as much as possible
// without exceeding original individual edge weights.
//
int main() {
int n;
cin >> n;
vector<node> nodes(n);
vector<edge> edges(n - 1);
for (int i = 0; i < n - 1; i++) {
cin >> edges[i];
nodes[edges[i].x].edges.push_back(i);
}
vector<edge> original = edges;
if (!lower(nodes, edges, 0)) {
cout << -1 << endl;
return 0;
}
raise(nodes, edges, original, 0, numeric_limits<int>::max());
cout << n << endl;
for (const edge& e : edges) {
cout << e << endl;
}
return 0;
}