|
@@ -386,6 +386,11 @@ func (pq *prioq) Add(car *Car) {
|
386
|
386
|
|
387
|
387
|
// Rebalance tree to respect invariant
|
388
|
388
|
var i = len(pq.bintree) - 1
|
|
389
|
+ pq.percolateUp(i)
|
|
390
|
+}
|
|
391
|
+
|
|
392
|
+func (pq *prioq) percolateUp(k int) {
|
|
393
|
+ var i = k
|
389
|
394
|
var p = (i - 1) / 2
|
390
|
395
|
for p >= 0 && pq.bintree[p].greater(pq.bintree[i]) {
|
391
|
396
|
pq.bintree[p], pq.bintree[i] = pq.bintree[i], pq.bintree[p]
|
|
@@ -396,25 +401,31 @@ func (pq *prioq) Add(car *Car) {
|
396
|
401
|
}
|
397
|
402
|
}
|
398
|
403
|
|
399
|
|
-func (pq *prioq) Pop() *Car {
|
|
404
|
+func (pq *prioq) RemoveAtIndex(k int) *Car {
|
400
|
405
|
if len(pq.bintree) == 0 {
|
401
|
406
|
return nil
|
402
|
407
|
}
|
403
|
408
|
|
404
|
|
- if len(pq.bintree) == 1 {
|
405
|
|
- elem := pq.bintree[0]
|
406
|
|
- pq.bintree = pq.bintree[:0]
|
|
409
|
+ if k == len(pq.bintree)-1 {
|
|
410
|
+ elem := pq.bintree[k]
|
|
411
|
+ pq.bintree = pq.bintree[:k]
|
407
|
412
|
return elem
|
408
|
413
|
}
|
409
|
414
|
|
410
|
|
- pq.bintree[0].pqindex = -1
|
411
|
|
- elem := pq.bintree[0]
|
412
|
|
- // Put last element at root
|
413
|
|
- pq.bintree[0] = pq.bintree[len(pq.bintree)-1]
|
414
|
|
- pq.bintree[0].pqindex = 0
|
|
415
|
+ pq.bintree[k].pqindex = -1
|
|
416
|
+ elem := pq.bintree[k]
|
|
417
|
+ // Put last element at hole
|
|
418
|
+ pq.bintree[k] = pq.bintree[len(pq.bintree)-1]
|
|
419
|
+ pq.bintree[k].pqindex = k
|
415
|
420
|
// Remove last element
|
416
|
421
|
pq.bintree = pq.bintree[:len(pq.bintree)-1]
|
417
|
422
|
|
|
423
|
+ // Oops! might need to go up
|
|
424
|
+ if k != 0 && !pq.bintree[k].greater(pq.bintree[(k-1)/2]) {
|
|
425
|
+ pq.percolateUp(k)
|
|
426
|
+ return elem
|
|
427
|
+ }
|
|
428
|
+
|
418
|
429
|
// 1 9
|
419
|
430
|
// 10 9 10 12
|
420
|
431
|
// 11 12 13 14 -> 11 12 13 14
|
|
@@ -422,7 +433,8 @@ func (pq *prioq) Pop() *Car {
|
422
|
433
|
|
423
|
434
|
// Rebalance tree to respect invariant
|
424
|
435
|
len := len(pq.bintree)
|
425
|
|
- i, left, right := 0, 0, 0
|
|
436
|
+ i := k
|
|
437
|
+ left, right := 0, 0
|
426
|
438
|
for {
|
427
|
439
|
left = 2*i + 1
|
428
|
440
|
right = 2*i + 2
|
|
@@ -463,18 +475,8 @@ func (pq *prioq) Pop() *Car {
|
463
|
475
|
return elem
|
464
|
476
|
}
|
465
|
477
|
|
466
|
|
-func (pq *prioq) RemoveAtIndex(k int) *Car {
|
467
|
|
- if k == 0 {
|
468
|
|
- return pq.Pop()
|
469
|
|
- }
|
470
|
|
- pq.bintree[k].pqindex = -1
|
471
|
|
- elem := pq.bintree[k]
|
472
|
|
- reassign := pq.bintree[k+1:]
|
473
|
|
- pq.bintree = pq.bintree[:k]
|
474
|
|
- for _, c := range reassign {
|
475
|
|
- pq.Add(c)
|
476
|
|
- }
|
477
|
|
- return elem
|
|
478
|
+func (pq *prioq) Pop() *Car {
|
|
479
|
+ return pq.RemoveAtIndex(0)
|
478
|
480
|
}
|
479
|
481
|
|
480
|
482
|
func (pq *prioq) empty() bool {
|