-
Notifications
You must be signed in to change notification settings - Fork 0
/
beam.go
61 lines (54 loc) · 1.31 KB
/
beam.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
package optimization
import (
"github.com/jakecoffman/graph/ds"
"time"
)
func greater(a, b int) bool {
return a > b
}
// Beam is like BFS but restricts the search space to save time by only looking at the best nodes.
func Beam[T GameState[T, U], U any](start T, beamSize int, limit time.Duration) []U {
beam := ds.NewPriorityQueue[T](greater)
beam.Push(start, 0)
nextStates := ds.NewPriorityQueue[T](greater)
startTime := time.Now()
best := start
for time.Now().Sub(startTime) < limit {
nextStates.Clear()
// beam size is restricted
for b := 0; b < beamSize; b++ {
if beam.Empty() {
break
}
current := beam.Pop()
moves := current.PossibleNextMoves()
if len(moves) == 0 {
// terminal state
if best.Evaluation() < current.Evaluation() {
best = current
}
continue
}
for _, move := range moves {
next := current.Apply(move)
nextStates.Push(next, next.Evaluation())
}
}
beam, nextStates = nextStates, beam
}
if best == start {
// in case we saw no terminal state, use the best so far
best = beam.Pop()
}
var path []U
current := best
for current != start {
path = append(path, current.CreatedBy())
current = current.CameFrom()
}
// reverse
for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
path[i], path[j] = path[j], path[i]
}
return path
}