-
Notifications
You must be signed in to change notification settings - Fork 7
/
16.kt
59 lines (54 loc) · 2.43 KB
/
16.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
import java.util.PriorityQueue
data class State(
var time: Int,
var current: String,
var elTime: Int? = null,
var elephant: String? = null,
var opened: Set<String> = setOf(),
var flow: Int = 0,
) : Comparable<State> {
override fun compareTo(other: State) = compareValuesBy(this, other, { -it.flow })
}
fun main() {
val input = generateSequence(::readlnOrNull).toList()
.map { Regex("([A-Z]{2}|\\d+)").findAll(it).toList().map { it.value } }
val neighbors = input.associate { it[0] to it.slice(2..it.size-1) }
val flows = input.associate { it[0] to it[1].toInt() }
fun getNonZeroNeighbors(curr: String, dist: Int = 0, visited: Set<String> = setOf()): Map<String, Int> {
val neigh = HashMap<String, Int>()
for (neighbor in neighbors[curr]!!.filter { it !in visited }) {
if (flows[neighbor]!! != 0)
neigh[neighbor] = dist+1
for ((name, d) in getNonZeroNeighbors(neighbor, dist+1, visited + setOf(curr)))
neigh[name] = minOf(d, neigh.getOrDefault(name, 1000))
}
return neigh
}
val nonZeroNeighbors = input.associate { it[0] to getNonZeroNeighbors(it[0]) }
fun solve(initialState: State): Int {
val queue = PriorityQueue<State>().also { it.add(initialState) }
var best = 0
val visited: MutableMap<List<String>, Int> = mutableMapOf()
while (queue.isNotEmpty()) {
var (time, current, elTime, elephant, opened, flow) = queue.remove()
best = maxOf(best, flow)
val vis = (opened.toList() + listOf(current, elephant ?: "")).sorted()
if (visited.getOrDefault(vis, -1) >= flow)
continue
visited[vis] = flow
if (elTime != null && elephant != null && time < elTime) {
time = elTime.also { elTime = time }
current = elephant.also { elephant = current }
}
for ((neighbor, dist) in nonZeroNeighbors[current]!!) {
val newTime = time-dist-1
val newFlow = flow+flows[neighbor]!!*newTime
if (newTime >= 0 && neighbor !in opened)
queue.add(State(newTime, neighbor, elTime, elephant, opened+setOf(neighbor), newFlow))
}
}
return best
}
solve(State(30, "AA")).run(::println)
solve(State(26, "AA", 26, "AA")).run(::println) // Takes ~10 seconds
}