浏览代码

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 年之前
父节点
当前提交
c51b556894
共有 1 个文件被更改,包括 24 次插入22 次删除
  1. 24 22
      main.go

+ 24 - 22
main.go

@@ -386,6 +386,11 @@ func (pq *prioq) Add(car *Car) {
386
 
386
 
387
 	// Rebalance tree to respect invariant
387
 	// Rebalance tree to respect invariant
388
 	var i = len(pq.bintree) - 1
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
 	var p = (i - 1) / 2
394
 	var p = (i - 1) / 2
390
 	for p >= 0 && pq.bintree[p].greater(pq.bintree[i]) {
395
 	for p >= 0 && pq.bintree[p].greater(pq.bintree[i]) {
391
 		pq.bintree[p], pq.bintree[i] = pq.bintree[i], pq.bintree[p]
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
 	if len(pq.bintree) == 0 {
405
 	if len(pq.bintree) == 0 {
401
 		return nil
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
 		return elem
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
 	// Remove last element
420
 	// Remove last element
416
 	pq.bintree = pq.bintree[:len(pq.bintree)-1]
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
 	//        1                  9
429
 	//        1                  9
419
 	//    10     9	         10     12
430
 	//    10     9	         10     12
420
 	//  11 12   13 14  ->  11 12   13 14
431
 	//  11 12   13 14  ->  11 12   13 14
@@ -422,7 +433,8 @@ func (pq *prioq) Pop() *Car {
422
 
433
 
423
 	// Rebalance tree to respect invariant
434
 	// Rebalance tree to respect invariant
424
 	len := len(pq.bintree)
435
 	len := len(pq.bintree)
425
-	i, left, right := 0, 0, 0
436
+	i := k
437
+	left, right := 0, 0
426
 	for {
438
 	for {
427
 		left = 2*i + 1
439
 		left = 2*i + 1
428
 		right = 2*i + 2
440
 		right = 2*i + 2
@@ -463,18 +475,8 @@ func (pq *prioq) Pop() *Car {
463
 	return elem
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
 func (pq *prioq) empty() bool {
482
 func (pq *prioq) empty() bool {