gitee.com/h79/goutils@v1.22.10/common/queue/priority.go (about)

     1  package queue
     2  
     3  import "time"
     4  
     5  // Priority 优先级
     6  type Priority struct {
     7  	l        []IPriority
     8  	location *time.Location
     9  }
    10  
    11  var _ Queue = (*Priority)(nil)
    12  
    13  func NewPriority() *Priority {
    14  	return &Priority{location: time.Local}
    15  }
    16  
    17  func (pq *Priority) SetLocation(l *time.Location) {
    18  	pq.location = l
    19  }
    20  
    21  // Len returns the Priority length.
    22  func (pq *Priority) Len() int { return len(pq.l) }
    23  
    24  func (pq *Priority) Add(elem any) error {
    25  	return pq.Push(elem)
    26  }
    27  
    28  func (pq *Priority) Push(x any) error {
    29  	p := NewPriorityItem(x, WithUnixNano(pq.location))
    30  	index := pq.getInsertIndex(p.PValue(), 0, pq.Len()-1)
    31  	pq.l = append(pq.l, nil)
    32  	copy(pq.l[index+1:], pq.l[index:])
    33  
    34  	pq.l[index] = p
    35  	return nil
    36  }
    37  
    38  // Pop implements the heap.Interface.Pop.
    39  // Removes and returns element Len() - 1.
    40  func (pq *Priority) Pop() any {
    41  	//old := pq.l
    42  	//n := len(old)
    43  	//if n == 0 {
    44  	//	return nil
    45  	//}
    46  	//item := old[n-1]
    47  	//pq.l = old[0 : n-1]
    48  	old := pq.l
    49  	item := old[0]
    50  	pq.l = old[1:]
    51  	return item.GValue()
    52  }
    53  
    54  // Head returns the first item of a Priority without removing it.
    55  func (pq *Priority) Head() IPriority {
    56  	if pq.Empty() {
    57  		return nil
    58  	}
    59  	return pq.l[0]
    60  }
    61  
    62  func (pq *Priority) Peek() any {
    63  	if pq.Empty() {
    64  		return nil
    65  	}
    66  	return pq.l[0].GValue()
    67  }
    68  
    69  // Remove removes and returns the element at index i from the Priority.
    70  func (pq *Priority) Remove(i int) any {
    71  	if i < 0 && i >= pq.Len() {
    72  		return nil
    73  	}
    74  	item := pq.l[i]
    75  	pq.l = append(pq.l[:i], pq.l[i+1:]...)
    76  	return item.GValue()
    77  }
    78  
    79  func (pq *Priority) Empty() bool {
    80  	return pq.Len() == 0
    81  }
    82  
    83  // Find return index -1 not found.
    84  func (pq *Priority) Find(fn IndexFunc) (any, int) {
    85  	for i := range pq.l {
    86  		if fn(pq.l[i], i) {
    87  			return pq.l[i], i
    88  		}
    89  	}
    90  	return nil, -1
    91  }
    92  
    93  func (pq *Priority) Foreach(fn IndexFunc) {
    94  	for i := range pq.l {
    95  		if fn(pq.l[i], i) {
    96  			return
    97  		}
    98  	}
    99  }
   100  
   101  func (pq *Priority) getInsertIndex(priority int64, leftIndex, rightIndex int) (index int) {
   102  	if len(pq.l) == 0 {
   103  		// 如果当前优先级切片没有元素，则插入的index就是0
   104  		return 0
   105  	}
   106  
   107  	length := rightIndex - leftIndex
   108  	if pq.l[leftIndex].PValue() >= priority {
   109  		// 如果当前切片中最小的元素都超过了插入的优先级，则插入位置应该是最左边
   110  		return leftIndex
   111  	}
   112  
   113  	if pq.l[rightIndex].PValue() <= priority {
   114  		// 如果当前切片中最大的元素都没超过插入的优先级，则插入位置应该是最右边
   115  		return rightIndex + 1
   116  	}
   117  
   118  	if length == 1 && pq.l[leftIndex].PValue() < priority && pq.l[rightIndex].PValue() >= priority {
   119  		// 如果插入的优先级刚好在仅有的两个优先级之间，则中间的位置就是插入位置
   120  		return leftIndex + 1
   121  	}
   122  
   123  	middleVal := pq.l[leftIndex+length/2].PValue()
   124  
   125  	// 这里用二分法递归的方式，一直寻找正确的插入位置
   126  	if priority <= middleVal {
   127  		return pq.getInsertIndex(priority, leftIndex, leftIndex+length/2)
   128  	}
   129  	return pq.getInsertIndex(priority, leftIndex+length/2, rightIndex)
   130  }