main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"math/rand"
 	"os"
 	"sort"
 )
@@ -65,6 +66,7 @@ func (rs ByEndtime) Less(i, j int) bool {
 
 type Scheduler interface {
 	Add(*Car)
+	RemoveAtIndex(k int) *Car
 	Pop() *Car
 }
 
@@ -76,7 +78,8 @@ type Car struct {
 	X       int
 	Y       int
 
-	score int
+	score   int
+	pqindex int // used for removal in prority queue
 }
 
 func (c *Car) Update(r *Ride) {
@@ -132,7 +135,23 @@ func max(a, b int) int {
 	}
 }
 
-func Choose(c *Car) *Ride {
+var bestTotalScore int
+
+// invariant : when entering and leaving function,
+// Sched has the "same" heap. Same means successive
+// popped values will be the same (but the internal
+// ordering of nodes may be different)
+func Choose(cumulativeScore int, depth int) {
+	if cumulativeScore > bestTotalScore {
+		bestTotalScore = cumulativeScore
+	}
+	c := Sched.Pop()
+	if c == nil {
+		// Stop recursion, empty car list,
+		// could not match cars to rides
+		return
+	}
+
 	var bestRides []struct {
 		r         *Ride
 		lenOfRide int
@@ -143,9 +162,9 @@ func Choose(c *Car) *Ride {
 		if r.used {
 			continue
 		}
-		// if r.Length() > 6000 {
-		// 	continue
-		// }
+		if r.Length() > 6000 {
+			continue
+		}
 		if r.f < c.EarliestFinish(r) {
 			continue
 		}
@@ -161,7 +180,10 @@ func Choose(c *Car) *Ride {
 			}{r, lenOfRide, total})
 			// shitty sort-of-correct-but-quite-incorrect way
 			// of picking next best n rides
-			n := 3
+			n := 1
+			if rand.Intn(max(1, int(N/5))) == 0 {
+				n = 3
+			}
 			if len(bestRides) > n {
 				bestRides = bestRides[len(bestRides)-n:]
 			}
@@ -171,31 +193,31 @@ func Choose(c *Car) *Ride {
 	// 	fmt.Printf("Picking %d %d -> %d %d\n", bestRide.a, bestRide.b, bestRide.x, bestRide.y)
 	// }
 	if len(bestRides) != 0 {
-		// pick the not quite best ride...hey some
-		// of the inputs already  give a better score !
-		return bestRides[0].r
-		//return bestRides[len(bestRides)-1].r
-	}
-	return nil
-}
-
-func assign() bool {
-	c := Sched.Pop()
-	if c == nil {
-		return false
-	}
-	r := Choose(c)
-	if r == nil {
-		return true
+		for _, br := range bestRides {
+			r := br.r
+			r.used = true
+			oldC := *c
+			c.Update(r)
+			Sched.Add(c)
+			// recursion 101
+			Choose(cumulativeScore+c.score-oldC.score, depth+1)
+			Sched.RemoveAtIndex(c.pqindex)
+			// pqindex is meaningless now but it's ok
+			// it will be fixed when c is added again
+			*c = oldC
+			r.used = false
+		}
+	} else {
+		// another car may still have other rides
+		Choose(cumulativeScore, depth)
 	}
-	r.used = true
-	c.Update(r)
+	// add back the one we popped at beginning of function
 	Sched.Add(c)
-
-	return true
 }
 
 func solve() {
+	rand.Seed(1)
+
 	sort.Sort(ByEndtime(Rides))
 
 	Sched = &prioq{}
@@ -212,19 +234,17 @@ func solve() {
 		Sched.Add(c)
 	}
 
-	for assign() {
-	}
+	// start recursion
+	Choose(0, 0)
 
-	totalScore := 0
 	for _, c := range Cars {
 		fmt.Fprintf(output, "%d", len(c.Rides))
 		for _, ri := range c.Rides {
 			fmt.Fprintf(output, " %d", ri)
 		}
 		fmt.Fprintf(output, "\n")
-		totalScore += c.score
 	}
-	fmt.Printf("%d\n", totalScore)
+	fmt.Printf("%d\n", bestTotalScore)
 }
 
 func main() {
@@ -295,31 +315,36 @@ type prioq struct {
 
 func (pq *prioq) Add(car *Car) {
 	pq.bintree = append(pq.bintree, car)
+	pq.bintree[len(pq.bintree)-1].pqindex = len(pq.bintree) - 1
 
 	// Rebalance tree to respect invariant
 	var i = len(pq.bintree) - 1
 	var p = (i - 1) / 2
 	for p >= 0 && pq.bintree[p].Arrival > pq.bintree[i].Arrival {
 		pq.bintree[p], pq.bintree[i] = pq.bintree[i], pq.bintree[p]
+		pq.bintree[p].pqindex = p
+		pq.bintree[i].pqindex = i
 		i = p
 		p = (i - 1) / 2
 	}
 }
 
-func (pq *prioq) Pop() *Car {
+func (pq *prioq) RemoveAtIndex(k int) *Car {
 	if len(pq.bintree) == 0 {
 		return nil
 	}
 
-	if len(pq.bintree) == 1 {
-		elem := pq.bintree[0]
-		pq.bintree = pq.bintree[:0]
+	if k == len(pq.bintree)-1 {
+		elem := pq.bintree[k]
+		pq.bintree = pq.bintree[:k]
 		return elem
 	}
 
-	elem := pq.bintree[0]
-	// Put last element at root
-	pq.bintree[0] = pq.bintree[len(pq.bintree)-1]
+	pq.bintree[k].pqindex = -1
+	elem := pq.bintree[k]
+	// Put last element at hole
+	pq.bintree[k] = pq.bintree[len(pq.bintree)-1]
+	pq.bintree[k].pqindex = k
 	// Remove last element
 	pq.bintree = pq.bintree[:len(pq.bintree)-1]
 
@@ -330,7 +355,8 @@ func (pq *prioq) Pop() *Car {
 
 	// Rebalance tree to respect invariant
 	len := len(pq.bintree)
-	i, left, right := 0, 0, 0
+	i := k
+	left, right := 0, 0
 	for {
 		left = 2*i + 1
 		right = 2*i + 2
@@ -340,6 +366,8 @@ func (pq *prioq) Pop() *Car {
 					break // Inferior to both children
 				} else {
 					pq.bintree[i], pq.bintree[left] = pq.bintree[left], pq.bintree[i]
+					pq.bintree[i].pqindex = i
+					pq.bintree[left].pqindex = left
 					i = left
 				}
 			} else {
@@ -347,6 +375,8 @@ func (pq *prioq) Pop() *Car {
 					break // Inferior to both children
 				} else {
 					pq.bintree[i], pq.bintree[right] = pq.bintree[right], pq.bintree[i]
+					pq.bintree[i].pqindex = i
+					pq.bintree[right].pqindex = right
 					i = right
 				}
 			}
@@ -355,6 +385,8 @@ func (pq *prioq) Pop() *Car {
 				break // Inferior to only child
 			}
 			pq.bintree[i], pq.bintree[left] = pq.bintree[left], pq.bintree[i]
+			pq.bintree[i].pqindex = i
+			pq.bintree[left].pqindex = left
 			i = left
 		} else { // No child
 			break
@@ -365,6 +397,10 @@ func (pq *prioq) Pop() *Car {
 	return elem
 }
 
+func (pq *prioq) Pop() *Car {
+	return pq.RemoveAtIndex(0)
+}
+
 func (pq *prioq) empty() bool {
 	return len(pq.bintree) == 0
 }