89 lines
1.9 KiB
Go
89 lines
1.9 KiB
Go
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
package bounded_tracking
|
|
|
|
import (
|
|
"container/heap"
|
|
|
|
"smariot.com/tsp/internal/solver/problem"
|
|
)
|
|
|
|
type minHeap[P problem.Problem[State], State comparable] struct {
|
|
problem problem.Problem[State]
|
|
known map[State]int
|
|
items []State
|
|
}
|
|
|
|
func (h minHeap[P, State]) Len() int {
|
|
return len(h.items)
|
|
}
|
|
|
|
func (h minHeap[P, State]) Less(i, j int) bool {
|
|
return h.problem.OptimisticLess(h.items[i], h.items[j])
|
|
}
|
|
|
|
func (h minHeap[P, State]) Swap(i, j int) {
|
|
h.items[i], h.items[j] = h.items[j], h.items[i]
|
|
h.known[h.items[i]] = i
|
|
h.known[h.items[j]] = j
|
|
}
|
|
|
|
func (h *minHeap[P, State]) Push(x any) {
|
|
state := x.(State)
|
|
h.items = append(h.items, state)
|
|
h.known[state] = len(h.items)
|
|
}
|
|
|
|
func (h *minHeap[P, State]) Pop() any {
|
|
n := len(h.items)
|
|
state := h.items[n-1]
|
|
h.items = h.items[:n-1]
|
|
delete(h.known, state)
|
|
return state
|
|
}
|
|
|
|
type solver[P problem.Problem[State], State comparable] struct {
|
|
minHeap[P, State]
|
|
}
|
|
|
|
func (s *solver[P, State]) Push(state State) {
|
|
if i, ok := s.known[state]; ok {
|
|
// The state is already in the heap, update its position instead.
|
|
heap.Fix(&s.minHeap, i)
|
|
return
|
|
}
|
|
|
|
heap.Push(&s.minHeap, state)
|
|
}
|
|
|
|
func (s *solver[P, State]) Pop() (State, bool) {
|
|
if s.Len() == 0 {
|
|
var zero State
|
|
return zero, false
|
|
}
|
|
|
|
return heap.Pop(&s.minHeap).(State), true
|
|
}
|
|
|
|
func (s *solver[P, State]) Reset() {
|
|
for _, state := range s.items {
|
|
s.problem.Discard(state)
|
|
}
|
|
|
|
s.items = s.items[:0]
|
|
clear(s.known)
|
|
}
|
|
|
|
// Returns a solver that for bounded problems that can update their states.
|
|
//
|
|
// Submitting a state that is already in the heap will update its position,
|
|
// rather than adding it again. problem.Discard will only be invoked once.
|
|
func New[P problem.Problem[State], State comparable](problem P) *solver[P, State] {
|
|
return &solver[P, State]{
|
|
minHeap: minHeap[P, State]{
|
|
problem: problem,
|
|
known: make(map[State]int),
|
|
},
|
|
}
|
|
}
|