/** * For each farm and for each day we want to find the minimum cost * to first arrive at that farm on that day. Given that we must * send a horse out on each successive day after first arriving at * a farm, we can solve this using a Minimum Cost Max Flow algorithm * if we know the cost to visit each later farm on each day. * * This submission gets WA due to a mismatch between the returned * cost for solving a subproblem and the cached cost for the subproblem. * * Time complexity: O(n^4) * * @author Finn Lidbetter */ import kotlin.math.* fun main() { val n = readln().toInt() var c = LongArray(n+1) var d = LongArray(n+1) var r = IntArray(n+1) for (i in 0..n-1) { val tokens = readln().split(" ").map{ it -> it.toInt() } c[i+1] = tokens[0].toLong() d[i+1] = tokens[1].toLong() r[i+1] = tokens[2] } var nodes = Array(n+1) { index -> Node(index, c[index], d[index], ArrayList()) } for (index in 1..n) { val parent = r[index] nodes[parent].children.add(index) } // cost[i][j] is the smallest cost to visit farmstead i on day j. var cost = Array(n+1) { LongArray(n+1) } // vis[i][j] is True iff we have computed the minimum cost of visiting farmstead i on day j. var vis = Array(n+1) { BooleanArray(n+1) } println(getMinCost(0, 0, nodes, cost, vis)) // var index = 0 // for (row in cost) { // println("Farm $index") // println(row.contentToString()) // index++ // } } fun getMinCost(farm: Int, day: Int, nodes: Array, precomputedCost: Array, isPrecomputed: Array): Long { if (isPrecomputed[farm][day]) { return precomputedCost[farm][day] } val nChildren = nodes[farm].children.size val dayDelta = nodes[farm].d - day.toLong() val selfCost = nodes[farm].c * dayDelta * dayDelta if (nChildren == 0) { precomputedCost[farm][day] = selfCost isPrecomputed[farm][day] = true return precomputedCost[farm][day] } var costMatrix = Array(nChildren) { childIndex -> LongArray(nChildren) { dayOffset -> getMinCost( nodes[farm].children.get(childIndex), day + dayOffset + 1, nodes, precomputedCost, isPrecomputed ) } } var numFlowNodes = 2 * nChildren + 2 var flowGraphCapacities = Array(numFlowNodes) { LongArray(numFlowNodes) } var flowGraphCosts = Array(numFlowNodes) { LongArray(numFlowNodes) } val flowSource = numFlowNodes - 2 val flowSink = numFlowNodes - 1 for (i in 0..nChildren - 1) { flowGraphCapacities[flowSource][i] = 1L flowGraphCapacities[nChildren + i][flowSink] = 1L for (j in 0..nChildren - 1) { flowGraphCapacities[i][nChildren + j] = 1L flowGraphCosts[i][nChildren + j] = costMatrix[i][j] } } val (flow, flowCost) = getFlowCost(flowGraphCapacities, flowGraphCosts, flowSource, flowSink) check(flow.toInt() == nChildren) { "Somehow the expected flow was not realised! Expected ${nChildren}, got ${flow}" } precomputedCost[farm][day] = flowCost // This is the bug! It should have + selfCost isPrecomputed[farm][day] = true return flowCost + selfCost } class Node(val id: Int, val c: Long, val d: Long, var children: ArrayList) fun getFlowCost(capacities: Array, costs: Array, source: Int, sink: Int): Pair { val numNodes = capacities.size var dad = IntArray(numNodes) var flow = Array(numNodes) { LongArray(numNodes) } var pi = LongArray(numNodes) var totalFlow = 0L var totalCost = 0L while (search(capacities, costs, source, sink, dad, flow, pi)) { var amt = Long.MAX_VALUE / 2L - 1L var x = sink while (x != source) { var otherAmt = flow[x][dad[x]] if (otherAmt == 0L) { otherAmt = capacities[dad[x]][x] - flow[dad[x]][x] } if (otherAmt < amt) { amt = otherAmt } x = dad[x] } x = sink while (x != source) { if (flow[x][dad[x]] != 0L) { flow[x][dad[x]] -= amt totalCost -= amt * costs[x][dad[x]] } else { flow[dad[x]][x] += amt totalCost += amt * costs[dad[x]][x] } x = dad[x] } totalFlow += amt } return Pair(totalFlow, totalCost) } fun search( capacities: Array, costs: Array, source: Int, sink: Int, dad: IntArray, flow: Array, pi: LongArray ): Boolean { val numNodes = capacities.size var found = BooleanArray(numNodes) var dist = LongArray(numNodes + 1) { Long.MAX_VALUE / 2L - 1L } dist[source] = 0L var curr = source while (curr != numNodes) { var best = numNodes found[curr] = true for (k in 0..numNodes - 1) { if (found[k]) { continue } if (flow[k][curr] != 0L) { val costSum = dist[curr] + pi[curr] - pi[k] - costs[k][curr] if (dist[k] > costSum) { dist[k] = costSum dad[k] = curr } } if (flow[curr][k] < capacities[curr][k]) { val costSum = dist[curr] + pi[curr] - pi[k] + costs[curr][k] if (dist[k] > costSum) { dist[k] = costSum dad[k] = curr } } if (dist[k] < dist[best]) { best = k } } curr = best } for (k in 0..numNodes - 1) { pi[k] = min(pi[k] + dist[k], Long.MAX_VALUE / 2L - 1L) } return found[sink] }