Prioq can now remove any element if its position is known

main.go

@@ -65,6 +65,7 @@ func (rs ByEndtime) Less(i, j int) bool {
 
 type Scheduler interface {
 	Add(*Car)
+	RemoveAtIndex(k int) *Car
 	Pop() *Car
 }
 
@@ -76,7 +77,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) {
@@ -295,31 +297,35 @@ 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]
+	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 +336,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 +347,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 +356,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 +366,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 +378,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
 }