-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathH.kt
77 lines (77 loc) · 2.25 KB
/
H.kt
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
fun main() {
val tests = readLine()!!.toInt()
repeat(tests) {
val vertices = readLine()!!.toInt()
val tree = Array(vertices) { arrayListOf<Pair<Int, Int>>() }
repeat(vertices - 1) {
val (v, u) = readLine()!!.split(' ').map { it.toInt() - 1 }
tree[v].add(u to it)
tree[u].add(v to it)
}
val answer = IntArray(vertices - 1)
val dp = IntArray(vertices)
fun paint(limit: Int): Boolean {
answer.fill(-1)
dp.fill(0)
var nextColor = 1
fun paint(vertex: Int, parent: Int, color: Int) {
for ((child, index) in tree[vertex]) if (child != parent) {
if (answer[index] == -1) {
answer[index] = color
paint(child, vertex, color)
}
}
}
fun dfs(vertex: Int, parent: Int = -1): Boolean {
val knapsack = IntArray(limit + 1) { -1 }
knapsack[0] = 0
var total = 0
for ((index, edge) in tree[vertex].withIndex()) if (edge.first != parent) {
val child = edge.first
if (!dfs(child, vertex)) {
return false
}
total += dp[child] + 1
for (w in limit - (dp[child] + 1) downTo 0) {
if (knapsack[w] >= 0 && knapsack[w + dp[child] + 1] < 0) {
knapsack[w + dp[child] + 1] = index
}
}
}
for (w in limit downTo 0) {
if (knapsack[w] >= 0) {
var cur = w
while (cur > 0) {
val (child, index) = tree[vertex][knapsack[cur]]
answer[index] = nextColor
paint(child, vertex, nextColor)
cur -= dp[child] + 1
}
dp[vertex] = total - w
if (w > 0) nextColor++
break
}
}
return dp[vertex] <= limit
}
if (dfs(0)) {
paint(0, -1, nextColor)
return true
}
return false
}
var left = 0
var right = vertices
while (right - left > 1) {
val test = (right + left) ushr 1
if (paint(test)) {
right = test
} else {
left = test
}
}
assert(paint(right))
println(right)
println(answer.joinToString(" "))
}
}