Browse Source

So there is a way to do prioq general node removal in log(n)

But it doesn't improve the overall speed of the algorithm much.
Meaning most of the time was not spent in that function.
Guillaume Koenig 7 years ago
parent
commit
c51b556894
1 changed files with 24 additions and 22 deletions
  1. 24 22
      main.go

+ 24 - 22
main.go

@@ -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 {