diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3fec32c --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +tmp/ diff --git a/cmd/solver-tune/main.go b/cmd/solver-tune/main.go new file mode 100644 index 0000000..8f12bf7 --- /dev/null +++ b/cmd/solver-tune/main.go @@ -0,0 +1,339 @@ +package main + +import ( + "encoding/json" + "flag" + "fmt" + "math/rand" + "os" + "slices" + "sort" + "strconv" + "strings" + "time" + + "rooms/solver" +) + +type tripData struct { + RoomSize int `json:"room_size"` + PreferNotMultiple int `json:"prefer_not_multiple"` + NoPreferCost int `json:"no_prefer_cost"` +} + +type studentData struct { + ID int64 `json:"id"` + Name string `json:"name"` +} + +type constraintsData struct { + Overalls []struct { + StudentAID int64 `json:"student_a_id"` + StudentBID int64 `json:"student_b_id"` + Kind string `json:"kind"` + } `json:"overalls"` +} + +func normalizeKey(a []int) string { + rm := map[int][]int{} + for i, room := range a { + rm[room] = append(rm[room], i) + } + var gs [][]int + for _, members := range rm { + slices.Sort(members) + gs = append(gs, members) + } + slices.SortFunc(gs, func(a, b []int) int { return a[0] - b[0] }) + var buf strings.Builder + for _, g := range gs { + for i, m := range g { + if i > 0 { + buf.WriteByte(',') + } + buf.WriteString(strconv.Itoa(m)) + } + buf.WriteByte(';') + } + return buf.String() +} + +type runResult struct { + score int + solutions [][]int + elapsed time.Duration +} + +func printStats(label string, results []runResult, runs int) { + scores := map[int]int{} + solutionSets := map[string]int{} + var totalTime time.Duration + var totalSolutions int + + for _, r := range results { + totalTime += r.elapsed + scores[r.score]++ + totalSolutions += len(r.solutions) + for _, sol := range r.solutions { + key := normalizeKey(sol) + solutionSets[key]++ + } + } + + fmt.Printf("--- %s ---\n", label) + fmt.Printf(" avg time: %v\n", totalTime/time.Duration(runs)) + + var scoreList []struct { + score int + count int + } + for s, c := range scores { + scoreList = append(scoreList, struct { + score int + count int + }{s, c}) + } + sort.Slice(scoreList, func(i, j int) bool { return scoreList[i].score > scoreList[j].score }) + + fmt.Printf(" score distribution:\n") + for _, sc := range scoreList { + fmt.Printf(" score %d: %d/%d runs (%.0f%%)\n", sc.score, sc.count, runs, float64(sc.count)/float64(runs)*100) + } + + fmt.Printf(" unique solutions seen: %d\n", len(solutionSets)) + fmt.Printf(" avg solutions per run: %.1f\n", float64(totalSolutions)/float64(runs)) + + var solFreqs []struct { + key string + count int + } + for k, c := range solutionSets { + solFreqs = append(solFreqs, struct { + key string + count int + }{k, c}) + } + sort.Slice(solFreqs, func(i, j int) bool { return solFreqs[i].count > solFreqs[j].count }) + + stableCount := 0 + for _, sf := range solFreqs { + if sf.count == runs { + stableCount++ + } + } + fmt.Printf(" solutions found in all runs: %d\n", stableCount) + if len(solFreqs) > 0 { + topN := min(5, len(solFreqs)) + fmt.Printf(" top %d solution frequencies: ", topN) + for i := range topN { + if i > 0 { + fmt.Print(", ") + } + fmt.Printf("%d/%d", solFreqs[i].count, runs) + } + fmt.Println() + } + fmt.Println() +} + +func main() { + dir := flag.String("dir", "tmp", "directory with trip/students/constraints JSON files") + runs := flag.Int("runs", 20, "number of solver runs per parameter set") + algo := flag.String("algo", "both", "algorithm: hillclimb, fast, sa, hybrid, both, all") + numRandom := flag.String("random", "30", "comma-separated random placement counts (hillclimb)") + numPerturb := flag.String("perturb", "200", "comma-separated perturbation counts (hillclimb)") + perturbMin := flag.Int("pmin", 2, "perturbation min groups (hillclimb)") + perturbMax := flag.Int("pmax", 5, "perturbation max groups (hillclimb)") + saRestarts := flag.String("restarts", "20", "comma-separated SA/hybrid restart counts") + saSteps := flag.String("steps", "10000", "comma-separated SA step counts") + hybridSteps := flag.String("hsteps", "5000", "comma-separated hybrid SA step counts") + saTempHigh := flag.Float64("thigh", 5.0, "SA initial temperature") + saTempLow := flag.Float64("tlow", 0.01, "SA final temperature") + hybridTempHigh := flag.Float64("hthigh", 10.0, "hybrid SA initial temperature") + hybridTempLow := flag.Float64("htlow", 0.1, "hybrid SA final temperature") + flag.Parse() + + tripBytes, err := os.ReadFile(*dir + "/1") + if err != nil { + fmt.Fprintf(os.Stderr, "reading trip: %v\n", err) + os.Exit(1) + } + var trip tripData + json.Unmarshal(tripBytes, &trip) + + studentsBytes, err := os.ReadFile(*dir + "/students") + if err != nil { + fmt.Fprintf(os.Stderr, "reading students: %v\n", err) + os.Exit(1) + } + var students []studentData + json.Unmarshal(studentsBytes, &students) + + constraintsBytes, err := os.ReadFile(*dir + "/constraints") + if err != nil { + fmt.Fprintf(os.Stderr, "reading constraints: %v\n", err) + os.Exit(1) + } + var cd constraintsData + json.Unmarshal(constraintsBytes, &cd) + + idx := map[int64]int{} + for i, s := range students { + idx[s.ID] = i + } + n := len(students) + + var constraints []solver.Constraint + for _, o := range cd.Overalls { + ai, aOk := idx[o.StudentAID] + bi, bOk := idx[o.StudentBID] + if !aOk || !bOk { + continue + } + constraints = append(constraints, solver.Constraint{ + StudentA: ai, + StudentB: bi, + Kind: o.Kind, + }) + } + + fmt.Printf("Students: %d, Room size: %d, Constraints: %d\n", n, trip.RoomSize, len(constraints)) + fmt.Printf("Prefer Not multiple: %d, No Prefer cost: %d\n", trip.PreferNotMultiple, trip.NoPreferCost) + fmt.Printf("Runs per config: %d\n\n", *runs) + + if *algo == "hillclimb" || *algo == "both" { + randomCounts := parseIntList(*numRandom) + perturbCounts := parseIntList(*numPerturb) + for _, nr := range randomCounts { + for _, np := range perturbCounts { + params := solver.Params{ + NumRandom: nr, + NumPerturb: np, + PerturbMin: *perturbMin, + PerturbMax: *perturbMax, + } + var results []runResult + for run := range *runs { + rng := rand.New(rand.NewSource(int64(run * 31337))) + start := time.Now() + sols := solver.Solve(n, trip.RoomSize, trip.PreferNotMultiple, trip.NoPreferCost, constraints, params, rng) + elapsed := time.Since(start) + if len(sols) > 0 { + var assignments [][]int + for _, s := range sols { + assignments = append(assignments, s.Assignment) + } + results = append(results, runResult{sols[0].Score, assignments, elapsed}) + } + } + label := fmt.Sprintf("hillclimb random=%d perturb=%d pmin=%d pmax=%d", nr, np, *perturbMin, *perturbMax) + printStats(label, results, *runs) + } + } + } + + if *algo == "fast" || *algo == "both" || *algo == "all" { + randomCounts := parseIntList(*numRandom) + perturbCounts := parseIntList(*numPerturb) + for _, nr := range randomCounts { + for _, np := range perturbCounts { + params := solver.Params{ + NumRandom: nr, + NumPerturb: np, + PerturbMin: *perturbMin, + PerturbMax: *perturbMax, + } + var results []runResult + for run := range *runs { + rng := rand.New(rand.NewSource(int64(run * 31337))) + start := time.Now() + sols := solver.SolveFast(n, trip.RoomSize, trip.PreferNotMultiple, trip.NoPreferCost, constraints, params, rng) + elapsed := time.Since(start) + if len(sols) > 0 { + var assignments [][]int + for _, s := range sols { + assignments = append(assignments, s.Assignment) + } + results = append(results, runResult{sols[0].Score, assignments, elapsed}) + } + } + label := fmt.Sprintf("fast random=%d perturb=%d pmin=%d pmax=%d", nr, np, *perturbMin, *perturbMax) + printStats(label, results, *runs) + } + } + } + + if *algo == "sa" || *algo == "all" { + restartCounts := parseIntList(*saRestarts) + stepCounts := parseIntList(*saSteps) + for _, nr := range restartCounts { + for _, ns := range stepCounts { + params := solver.SAParams{ + Restarts: nr, + Steps: ns, + TempHigh: *saTempHigh, + TempLow: *saTempLow, + } + var results []runResult + for run := range *runs { + rng := rand.New(rand.NewSource(int64(run * 31337))) + start := time.Now() + sols := solver.SolveSA(n, trip.RoomSize, trip.PreferNotMultiple, trip.NoPreferCost, constraints, params, rng) + elapsed := time.Since(start) + if len(sols) > 0 { + var assignments [][]int + for _, s := range sols { + assignments = append(assignments, s.Assignment) + } + results = append(results, runResult{sols[0].Score, assignments, elapsed}) + } + } + label := fmt.Sprintf("sa restarts=%d steps=%d thigh=%.1f tlow=%.3f", nr, ns, *saTempHigh, *saTempLow) + printStats(label, results, *runs) + } + } + } + + if *algo == "hybrid" || *algo == "both" || *algo == "all" { + restartCounts := parseIntList(*saRestarts) + stepCounts := parseIntList(*hybridSteps) + for _, nr := range restartCounts { + for _, ns := range stepCounts { + params := solver.HybridParams{ + SARestarts: nr, + SASteps: ns, + TempHigh: *hybridTempHigh, + TempLow: *hybridTempLow, + } + var results []runResult + for run := range *runs { + rng := rand.New(rand.NewSource(int64(run * 31337))) + start := time.Now() + sols := solver.SolveHybrid(n, trip.RoomSize, trip.PreferNotMultiple, trip.NoPreferCost, constraints, params, rng) + elapsed := time.Since(start) + if len(sols) > 0 { + var assignments [][]int + for _, s := range sols { + assignments = append(assignments, s.Assignment) + } + results = append(results, runResult{sols[0].Score, assignments, elapsed}) + } + } + label := fmt.Sprintf("hybrid restarts=%d steps=%d thigh=%.1f tlow=%.3f", nr, ns, *hybridTempHigh, *hybridTempLow) + printStats(label, results, *runs) + } + } + } +} + +func parseIntList(s string) []int { + parts := strings.Split(s, ",") + var result []int + for _, p := range parts { + v, err := strconv.Atoi(strings.TrimSpace(p)) + if err == nil { + result = append(result, v) + } + } + return result +} diff --git a/main.go b/main.go index e311dde..8fa8b03 100644 --- a/main.go +++ b/main.go @@ -21,6 +21,8 @@ import ( "github.com/lib/pq" "google.golang.org/api/idtoken" + + "rooms/solver" ) //go:embed schema.sql @@ -1167,13 +1169,13 @@ func handleSolve(db *sql.DB) http.HandlerFunc { } defer crows.Close() - type constraint struct { - aID, bID int64 - kind, level string + type dbConstraint struct { + aID, bID int64 + kind, level string } - var allConstraints []constraint + var allConstraints []dbConstraint for crows.Next() { - var c constraint + var c dbConstraint if err := crows.Scan(&c.aID, &c.bID, &c.kind, &c.level); err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return @@ -1182,7 +1184,6 @@ func handleSolve(db *sql.DB) http.HandlerFunc { } type pairKey struct{ a, b int64 } - overalls := map[pairKey]string{} byPair := map[pairKey]map[string]string{} for _, c := range allConstraints { pk := pairKey{c.aID, c.bID} @@ -1192,6 +1193,7 @@ func handleSolve(db *sql.DB) http.HandlerFunc { byPair[pk][c.level] = c.kind } levelPriority := []string{"admin", "parent", "student"} + overalls := map[pairKey]string{} for pk, levels := range byPair { for _, lev := range levelPriority { if kind, ok := levels[lev]; ok { @@ -1207,349 +1209,25 @@ func handleSolve(db *sql.DB) http.HandlerFunc { } n := len(studentIDs) - mustTogether := map[[2]int]bool{} - mustApart := map[[2]int]bool{} + var constraints []solver.Constraint for pk, kind := range overalls { - ai, bi := idx[pk.a], idx[pk.b] - switch kind { - case "must": - p := [2]int{ai, bi} - if p[0] > p[1] { p[0], p[1] = p[1], p[0] } - mustTogether[p] = true - case "must_not": - p := [2]int{ai, bi} - if p[0] > p[1] { p[0], p[1] = p[1], p[0] } - mustApart[p] = true - } + constraints = append(constraints, solver.Constraint{ + StudentA: idx[pk.a], + StudentB: idx[pk.b], + Kind: kind, + }) } - uf := make([]int, n) - for i := range uf { uf[i] = i } - var ufFind func(int) int - ufFind = func(x int) int { - if uf[x] != x { uf[x] = ufFind(uf[x]) } - return uf[x] - } - ufUnion := func(a, b int) { - ra, rb := ufFind(a), ufFind(b) - if ra != rb { uf[ra] = rb } - } + rng := rand.New(rand.NewSource(rand.Int63())) + solutions := solver.SolveFast(n, roomSize, pnMultiple, npCost, constraints, solver.DefaultParams, rng) - for p := range mustTogether { - ufUnion(p[0], p[1]) - } - - hasConflict := false - for p := range mustApart { - if ufFind(p[0]) == ufFind(p[1]) { - hasConflict = true - break - } - } - - if hasConflict { + if solutions == nil { http.Error(w, "hard conflicts exist, resolve before solving", http.StatusBadRequest) return } - groups := map[int][]int{} - for i := range n { - root := ufFind(i) - groups[root] = append(groups[root], i) - } - - hasPrefer := make([]bool, n) - for pk, kind := range overalls { - if kind == "prefer" { - hasPrefer[idx[pk.a]] = true - } - } - - score := func(assignment []int) int { - s := 0 - gotPrefer := make([]bool, n) - for pk, kind := range overalls { - ai, bi := idx[pk.a], idx[pk.b] - sameRoom := assignment[ai] == assignment[bi] - switch kind { - case "prefer": - if sameRoom { - s++ - gotPrefer[ai] = true - } - case "prefer_not": - if sameRoom { s -= pnMultiple } - } - } - for i := range n { - if hasPrefer[i] && !gotPrefer[i] { - s -= npCost - } - } - return s - } - - feasible := func(assignment []int) bool { - for p := range mustApart { - if assignment[p[0]] == assignment[p[1]] { return false } - } - roomCounts := map[int]int{} - for _, room := range assignment { - roomCounts[room]++ - } - for _, cnt := range roomCounts { - if cnt > roomSize { return false } - } - return true - } - numRooms := (n + roomSize - 1) / roomSize - assignment := make([]int, n) - groupList := make([][]int, 0, len(groups)) - for _, members := range groups { - groupList = append(groupList, members) - } - slices.SortFunc(groupList, func(a, b []int) int { return len(b) - len(a) }) - - roomCap := make([]int, numRooms) - for i := range roomCap { roomCap[i] = roomSize } - - placed := false - var placeGroups func(gi int) bool - placeGroups = func(gi int) bool { - if gi >= len(groupList) { return true } - grp := groupList[gi] - for room := range numRooms { - if roomCap[room] < len(grp) { continue } - ok := true - for _, member := range grp { - for p := range mustApart { - partner := -1 - if p[0] == member { partner = p[1] } - if p[1] == member { partner = p[0] } - if partner >= 0 && assignment[partner] == room { - alreadyPlaced := false - for gj := range gi { - if slices.Contains(groupList[gj], partner) { - alreadyPlaced = true - break - } - } - if alreadyPlaced { ok = false; break } - } - } - if !ok { break } - } - if !ok { continue } - for _, member := range grp { assignment[member] = room } - roomCap[room] -= len(grp) - if placeGroups(gi + 1) { return true } - roomCap[room] += len(grp) - } - return false - } - placed = placeGroups(0) - - if !placed { - for i := range n { - assignment[i] = i % numRooms - } - } - - initialAssignment := make([]int, n) - copy(initialAssignment, assignment) - - bestScore := score(assignment) - var bestSolutions [][]int - seen := map[string]bool{} - normalizeKey := func(a []int) string { - rm := map[int][]int{} - for i, room := range a { - rm[room] = append(rm[room], i) - } - var gs [][]int - for _, members := range rm { - slices.Sort(members) - gs = append(gs, members) - } - slices.SortFunc(gs, func(a, b []int) int { return a[0] - b[0] }) - var buf strings.Builder - for _, g := range gs { - for i, m := range g { - if i > 0 { - buf.WriteByte(',') - } - buf.WriteString(strconv.Itoa(m)) - } - buf.WriteByte(';') - } - return buf.String() - } - addSolution := func(a []int, s int) { - if s > bestScore { - bestScore = s - bestSolutions = nil - seen = map[string]bool{} - } - if s == bestScore { - key := normalizeKey(a) - if !seen[key] { - seen[key] = true - bestSolutions = append(bestSolutions, slices.Clone(a)) - } - } - } - addSolution(assignment, bestScore) - - roomCount := func(a []int, room int) int { - c := 0 - for _, r := range a { - if r == room { c++ } - } - return c - } - - uniqueGroups := make([]int, 0, len(groups)) - for root := range groups { - uniqueGroups = append(uniqueGroups, root) - } - slices.Sort(uniqueGroups) - - hillClimb := func(assignment []int) int { - currentScore := score(assignment) - for { - bestDelta := 0 - bestMove := -1 - bestTarget := -1 - bestSwapJ := -1 - - for gi, gRoot := range uniqueGroups { - grp := groups[gRoot] - gRoom := assignment[grp[0]] - - for room := range numRooms { - if room == gRoom { continue } - if roomCount(assignment, room)+len(grp) > roomSize { continue } - for _, m := range grp { assignment[m] = room } - if feasible(assignment) { - delta := score(assignment) - currentScore - if delta > bestDelta { - bestDelta = delta - bestMove = gi - bestTarget = room - bestSwapJ = -1 - } - } - for _, m := range grp { assignment[m] = gRoom } - } - - for gj := gi + 1; gj < len(uniqueGroups); gj++ { - grp2 := groups[uniqueGroups[gj]] - g2Room := assignment[grp2[0]] - if gRoom == g2Room { continue } - newGRoom := roomCount(assignment, gRoom) - len(grp) + len(grp2) - newG2Room := roomCount(assignment, g2Room) - len(grp2) + len(grp) - if newGRoom > roomSize || newG2Room > roomSize { continue } - for _, m := range grp { assignment[m] = g2Room } - for _, m := range grp2 { assignment[m] = gRoom } - if feasible(assignment) { - delta := score(assignment) - currentScore - if delta > bestDelta { - bestDelta = delta - bestMove = gi - bestTarget = -1 - bestSwapJ = gj - } - } - for _, m := range grp { assignment[m] = gRoom } - for _, m := range grp2 { assignment[m] = g2Room } - } - } - - if bestDelta <= 0 { break } - - grp := groups[uniqueGroups[bestMove]] - gRoom := assignment[grp[0]] - if bestSwapJ < 0 { - for _, m := range grp { assignment[m] = bestTarget } - } else { - grp2 := groups[uniqueGroups[bestSwapJ]] - g2Room := assignment[grp2[0]] - for _, m := range grp { assignment[m] = g2Room } - for _, m := range grp2 { assignment[m] = gRoom } - } - currentScore += bestDelta - } - return currentScore - } - - randomPlacement := func() bool { - perm := rand.Perm(len(groupList)) - for i := range roomCap { roomCap[i] = roomSize } - for _, pi := range perm { - grp := groupList[pi] - placed := false - order := rand.Perm(numRooms) - for _, room := range order { - if roomCap[room] < len(grp) { continue } - valid := true - for _, member := range grp { - for p := range mustApart { - partner := -1 - if p[0] == member { partner = p[1] } - if p[1] == member { partner = p[0] } - if partner >= 0 && assignment[partner] == room { - valid = false - break - } - } - if !valid { break } - } - if !valid { continue } - for _, member := range grp { assignment[member] = room } - roomCap[room] -= len(grp) - placed = true - break - } - if !placed { return false } - } - return true - } - - perturb := func(src []int, count int) { - copy(assignment, src) - indices := rand.Perm(len(uniqueGroups)) - count = min(count, len(indices)) - for _, gi := range indices[:count] { - grp := groups[uniqueGroups[gi]] - oldRoom := assignment[grp[0]] - rooms := rand.Perm(numRooms) - for _, room := range rooms { - if room == oldRoom { continue } - if roomCount(assignment, room)+len(grp) > roomSize { continue } - for _, m := range grp { assignment[m] = room } - if feasible(assignment) { break } - for _, m := range grp { assignment[m] = oldRoom } - } - } - } - - copy(assignment, initialAssignment) - addSolution(assignment, hillClimb(assignment)) - - for range 30 { - if randomPlacement() { - addSolution(assignment, hillClimb(assignment)) - } - } - - for range 200 { - src := bestSolutions[rand.Intn(len(bestSolutions))] - perturb(src, 2+rand.Intn(3)) - addSolution(assignment, hillClimb(assignment)) - } - type roomMember struct { ID int64 `json:"id"` Name string `json:"name"` @@ -1559,9 +1237,9 @@ func handleSolve(db *sql.DB) http.HandlerFunc { Score int `json:"score"` } var results []solutionResult - for _, sol := range bestSolutions { + for _, sol := range solutions { roomMap := map[int][]roomMember{} - for i, room := range sol { + for i, room := range sol.Assignment { sid := studentIDs[i] roomMap[room] = append(roomMap[room], roomMember{ID: sid, Name: studentName[sid]}) } @@ -1573,7 +1251,7 @@ func handleSolve(db *sql.DB) http.HandlerFunc { } } slices.SortFunc(rooms, func(a, b []roomMember) int { return strings.Compare(a[0].Name, b[0].Name) }) - results = append(results, solutionResult{Rooms: rooms, Score: bestScore}) + results = append(results, solutionResult{Rooms: rooms, Score: sol.Score}) } slices.SortFunc(results, func(a, b solutionResult) int { for i := range min(len(a.Rooms), len(b.Rooms)) { diff --git a/solver/solver.go b/solver/solver.go new file mode 100644 index 0000000..ea52b00 --- /dev/null +++ b/solver/solver.go @@ -0,0 +1,1373 @@ +package solver + +import ( + "math" + "math/rand" + "slices" + "strconv" + "strings" +) + +type Constraint struct { + StudentA int + StudentB int + Kind string +} + +type Params struct { + NumRandom int + NumPerturb int + PerturbMin int + PerturbMax int +} + +var DefaultParams = Params{ + NumRandom: 50, + NumPerturb: 750, + PerturbMin: 3, + PerturbMax: 8, +} + +type SAParams struct { + Restarts int + Steps int + TempHigh float64 + TempLow float64 +} + +var DefaultSAParams = SAParams{ + Restarts: 20, + Steps: 10000, + TempHigh: 5.0, + TempLow: 0.01, +} + +type HybridParams struct { + SARestarts int + SASteps int + TempHigh float64 + TempLow float64 +} + +var DefaultHybridParams = HybridParams{ + SARestarts: 50, + SASteps: 5000, + TempHigh: 10.0, + TempLow: 0.1, +} + +type Solution struct { + Assignment []int + Score int +} + +func normalizeKey(a []int) string { + rm := map[int][]int{} + for i, room := range a { + rm[room] = append(rm[room], i) + } + var gs [][]int + for _, members := range rm { + slices.Sort(members) + gs = append(gs, members) + } + slices.SortFunc(gs, func(a, b []int) int { return a[0] - b[0] }) + var buf strings.Builder + for _, g := range gs { + for i, m := range g { + if i > 0 { + buf.WriteByte(',') + } + buf.WriteString(strconv.Itoa(m)) + } + buf.WriteByte(';') + } + return buf.String() +} + +type solverState struct { + n int + roomSize int + numRooms int + pnMultiple int + npCost int + + constraints []Constraint + mustApart map[[2]int]bool + + groups map[int][]int + groupList [][]int + groupOf []int + uniqueGroups []int + + studentConstraints [][]int + hasPrefer []bool + preferFrom [][]int + mustApartFor [][]int +} + +func newSolverState(n, roomSize, pnMultiple, npCost int, constraints []Constraint) *solverState { + s := &solverState{ + n: n, + roomSize: roomSize, + numRooms: (n + roomSize - 1) / roomSize, + pnMultiple: pnMultiple, + npCost: npCost, + constraints: constraints, + mustApart: map[[2]int]bool{}, + } + + mustTogether := map[[2]int]bool{} + for _, c := range constraints { + switch c.Kind { + case "must": + p := [2]int{c.StudentA, c.StudentB} + if p[0] > p[1] { + p[0], p[1] = p[1], p[0] + } + mustTogether[p] = true + case "must_not": + p := [2]int{c.StudentA, c.StudentB} + if p[0] > p[1] { + p[0], p[1] = p[1], p[0] + } + s.mustApart[p] = true + } + } + + uf := make([]int, n) + for i := range uf { + uf[i] = i + } + var ufFind func(int) int + ufFind = func(x int) int { + if uf[x] != x { + uf[x] = ufFind(uf[x]) + } + return uf[x] + } + for p := range mustTogether { + ra, rb := ufFind(p[0]), ufFind(p[1]) + if ra != rb { + uf[ra] = rb + } + } + + s.groups = map[int][]int{} + for i := range n { + root := ufFind(i) + s.groups[root] = append(s.groups[root], i) + } + + s.groupOf = make([]int, n) + for root, members := range s.groups { + for _, m := range members { + s.groupOf[m] = root + } + } + + s.groupList = make([][]int, 0, len(s.groups)) + for _, members := range s.groups { + s.groupList = append(s.groupList, members) + } + slices.SortFunc(s.groupList, func(a, b []int) int { return len(b) - len(a) }) + + s.uniqueGroups = make([]int, 0, len(s.groups)) + for root := range s.groups { + s.uniqueGroups = append(s.uniqueGroups, root) + } + slices.Sort(s.uniqueGroups) + + s.studentConstraints = make([][]int, n) + s.hasPrefer = make([]bool, n) + s.preferFrom = make([][]int, n) + s.mustApartFor = make([][]int, n) + for ci, c := range constraints { + s.studentConstraints[c.StudentA] = append(s.studentConstraints[c.StudentA], ci) + s.studentConstraints[c.StudentB] = append(s.studentConstraints[c.StudentB], ci) + if c.Kind == "prefer" { + s.hasPrefer[c.StudentA] = true + s.preferFrom[c.StudentA] = append(s.preferFrom[c.StudentA], ci) + } + } + for p := range s.mustApart { + s.mustApartFor[p[0]] = append(s.mustApartFor[p[0]], p[1]) + s.mustApartFor[p[1]] = append(s.mustApartFor[p[1]], p[0]) + } + + return s +} + +func (s *solverState) hasHardConflict() bool { + uf := make([]int, s.n) + for i := range uf { + uf[i] = i + } + var ufFind func(int) int + ufFind = func(x int) int { + if uf[x] != x { + uf[x] = ufFind(uf[x]) + } + return uf[x] + } + for _, c := range s.constraints { + if c.Kind == "must" { + ra, rb := ufFind(c.StudentA), ufFind(c.StudentB) + if ra != rb { + uf[ra] = rb + } + } + } + for p := range s.mustApart { + if ufFind(p[0]) == ufFind(p[1]) { + return true + } + } + return false +} + +func (s *solverState) score(assignment []int) int { + sc := 0 + gotPrefer := make([]bool, s.n) + for _, c := range s.constraints { + sameRoom := assignment[c.StudentA] == assignment[c.StudentB] + switch c.Kind { + case "prefer": + if sameRoom { + sc++ + gotPrefer[c.StudentA] = true + } + case "prefer_not": + if sameRoom { + sc -= s.pnMultiple + } + } + } + for i := range s.n { + if s.hasPrefer[i] && !gotPrefer[i] { + sc -= s.npCost + } + } + return sc +} + +func (s *solverState) roomCounts(assignment []int) []int { + counts := make([]int, s.numRooms) + for _, room := range assignment { + counts[room]++ + } + return counts +} + +func (s *solverState) feasibleForGroup(assignment []int, groupRoot int, room int) bool { + for _, m := range s.groups[groupRoot] { + for _, partner := range s.mustApartFor[m] { + if s.groupOf[partner] != groupRoot && assignment[partner] == room { + return false + } + } + } + return true +} + +func (s *solverState) fastHillClimb(assignment []int) int { + n := s.n + roomSize := s.roomSize + numRooms := s.numRooms + + roomCounts := make([]int, numRooms) + for _, room := range assignment { + roomCounts[room]++ + } + + prefSatCount := make([]int, n) + for _, c := range s.constraints { + if c.Kind == "prefer" && assignment[c.StudentA] == assignment[c.StudentB] { + prefSatCount[c.StudentA]++ + } + } + + currentScore := 0 + for _, c := range s.constraints { + if assignment[c.StudentA] == assignment[c.StudentB] { + switch c.Kind { + case "prefer": + currentScore++ + case "prefer_not": + currentScore -= s.pnMultiple + } + } + } + for i := range n { + if s.hasPrefer[i] && prefSatCount[i] == 0 { + currentScore -= s.npCost + } + } + + memberSet := make([]bool, n) + + deltaForMove := func(groupRoot int, oldRoom, newRoom int) int { + members := s.groups[groupRoot] + for _, m := range members { + memberSet[m] = true + } + + delta := 0 + npAffected := make(map[int]int) + + for _, m := range members { + for _, ci := range s.studentConstraints[m] { + c := s.constraints[ci] + other := c.StudentB + if other == m { + other = c.StudentA + } + if memberSet[other] { + continue + } + otherRoom := assignment[other] + wasSame := otherRoom == oldRoom + willBeSame := otherRoom == newRoom + if wasSame == willBeSame { + continue + } + switch c.Kind { + case "prefer": + if wasSame { + delta-- + npAffected[c.StudentA]-- + } else { + delta++ + npAffected[c.StudentA]++ + } + case "prefer_not": + if wasSame { + delta += s.pnMultiple + } else { + delta -= s.pnMultiple + } + } + } + } + + for student, change := range npAffected { + if !s.hasPrefer[student] { + continue + } + wasSat := prefSatCount[student] > 0 + willBeSat := prefSatCount[student]+change > 0 + if wasSat && !willBeSat { + delta -= s.npCost + } else if !wasSat && willBeSat { + delta += s.npCost + } + } + + for _, m := range members { + memberSet[m] = false + } + return delta + } + + applyMove := func(groupRoot int, oldRoom, newRoom int) { + members := s.groups[groupRoot] + for _, m := range members { + memberSet[m] = true + } + for _, m := range members { + for _, ci := range s.studentConstraints[m] { + c := s.constraints[ci] + other := c.StudentB + if other == m { + other = c.StudentA + } + if memberSet[other] { + continue + } + otherRoom := assignment[other] + wasSame := otherRoom == oldRoom + willBeSame := otherRoom == newRoom + if wasSame == willBeSame { + continue + } + if c.Kind == "prefer" { + if wasSame { + prefSatCount[c.StudentA]-- + } else { + prefSatCount[c.StudentA]++ + } + } + } + } + roomCounts[oldRoom] -= len(members) + for _, m := range members { + assignment[m] = newRoom + } + roomCounts[newRoom] += len(members) + for _, m := range members { + memberSet[m] = false + } + } + + for { + bestDelta := 0 + bestGi := -1 + bestTarget := -1 + bestSwapGj := -1 + + for gi, gRoot := range s.uniqueGroups { + grp := s.groups[gRoot] + gRoom := assignment[grp[0]] + + for room := range numRooms { + if room == gRoom { + continue + } + if roomCounts[room]+len(grp) > roomSize { + continue + } + if !s.feasibleForGroup(assignment, gRoot, room) { + continue + } + delta := deltaForMove(gRoot, gRoom, room) + if delta > bestDelta { + bestDelta = delta + bestGi = gi + bestTarget = room + bestSwapGj = -1 + } + } + + for gj := gi + 1; gj < len(s.uniqueGroups); gj++ { + g2Root := s.uniqueGroups[gj] + grp2 := s.groups[g2Root] + g2Room := assignment[grp2[0]] + if gRoom == g2Room { + continue + } + newGRoom := roomCounts[gRoom] - len(grp) + len(grp2) + newG2Room := roomCounts[g2Room] - len(grp2) + len(grp) + if newGRoom > roomSize || newG2Room > roomSize { + continue + } + if !s.feasibleForGroup(assignment, gRoot, g2Room) { + continue + } + delta1 := deltaForMove(gRoot, gRoom, g2Room) + applyMove(gRoot, gRoom, g2Room) + if !s.feasibleForGroup(assignment, g2Root, gRoom) { + applyMove(gRoot, g2Room, gRoom) + continue + } + delta2 := deltaForMove(g2Root, g2Room, gRoom) + applyMove(gRoot, g2Room, gRoom) + + totalDelta := delta1 + delta2 + if totalDelta > bestDelta { + bestDelta = totalDelta + bestGi = gi + bestTarget = -1 + bestSwapGj = gj + } + } + } + + if bestDelta <= 0 { + break + } + + gRoot := s.uniqueGroups[bestGi] + gRoom := assignment[s.groups[gRoot][0]] + if bestSwapGj < 0 { + applyMove(gRoot, gRoom, bestTarget) + } else { + g2Root := s.uniqueGroups[bestSwapGj] + g2Room := assignment[s.groups[g2Root][0]] + applyMove(gRoot, gRoom, g2Room) + applyMove(g2Root, g2Room, gRoom) + } + currentScore += bestDelta + } + return currentScore +} + +func (s *solverState) initialPlacement(assignment []int) bool { + roomCap := make([]int, s.numRooms) + for i := range roomCap { + roomCap[i] = s.roomSize + } + + var placeGroups func(gi int) bool + placeGroups = func(gi int) bool { + if gi >= len(s.groupList) { + return true + } + grp := s.groupList[gi] + for room := range s.numRooms { + if roomCap[room] < len(grp) { + continue + } + ok := true + for _, member := range grp { + for p := range s.mustApart { + partner := -1 + if p[0] == member { + partner = p[1] + } + if p[1] == member { + partner = p[0] + } + if partner >= 0 && assignment[partner] == room { + alreadyPlaced := false + for gj := range gi { + if slices.Contains(s.groupList[gj], partner) { + alreadyPlaced = true + break + } + } + if alreadyPlaced { + ok = false + break + } + } + } + if !ok { + break + } + } + if !ok { + continue + } + for _, member := range grp { + assignment[member] = room + } + roomCap[room] -= len(grp) + if placeGroups(gi + 1) { + return true + } + roomCap[room] += len(grp) + } + return false + } + return placeGroups(0) +} + +func (s *solverState) randomPlacement(assignment []int, rng *rand.Rand) bool { + roomCap := make([]int, s.numRooms) + for i := range roomCap { + roomCap[i] = s.roomSize + } + perm := rng.Perm(len(s.groupList)) + for _, pi := range perm { + grp := s.groupList[pi] + placed := false + order := rng.Perm(s.numRooms) + for _, room := range order { + if roomCap[room] < len(grp) { + continue + } + valid := true + for _, member := range grp { + for p := range s.mustApart { + partner := -1 + if p[0] == member { + partner = p[1] + } + if p[1] == member { + partner = p[0] + } + if partner >= 0 && assignment[partner] == room { + valid = false + break + } + } + if !valid { + break + } + } + if !valid { + continue + } + for _, member := range grp { + assignment[member] = room + } + roomCap[room] -= len(grp) + placed = true + break + } + if !placed { + return false + } + } + return true +} + +type solutionTracker struct { + bestScore int + bestSolutions [][]int + seen map[string]bool +} + +func newTracker(initial []int, score int) *solutionTracker { + t := &solutionTracker{ + bestScore: score, + seen: map[string]bool{}, + } + key := normalizeKey(initial) + t.seen[key] = true + t.bestSolutions = append(t.bestSolutions, slices.Clone(initial)) + return t +} + +func (t *solutionTracker) add(a []int, s int) { + if s > t.bestScore { + t.bestScore = s + t.bestSolutions = nil + t.seen = map[string]bool{} + } + if s == t.bestScore { + key := normalizeKey(a) + if !t.seen[key] { + t.seen[key] = true + t.bestSolutions = append(t.bestSolutions, slices.Clone(a)) + } + } +} + +func Solve(n, roomSize, pnMultiple, npCost int, constraints []Constraint, params Params, rng *rand.Rand) []Solution { + if n == 0 { + return nil + } + + st := newSolverState(n, roomSize, pnMultiple, npCost, constraints) + if st.hasHardConflict() { + return nil + } + + assignment := make([]int, n) + if !st.initialPlacement(assignment) { + for i := range n { + assignment[i] = i % st.numRooms + } + } + + initialAssignment := slices.Clone(assignment) + tracker := newTracker(assignment, st.score(assignment)) + + roomCount := func(a []int, room int) int { + c := 0 + for _, r := range a { + if r == room { + c++ + } + } + return c + } + + feasible := func(a []int) bool { + for p := range st.mustApart { + if a[p[0]] == a[p[1]] { + return false + } + } + rc := map[int]int{} + for _, room := range a { + rc[room]++ + } + for _, cnt := range rc { + if cnt > roomSize { + return false + } + } + return true + } + + hillClimb := func(assignment []int) int { + currentScore := st.score(assignment) + for { + bestDelta := 0 + bestMove := -1 + bestTarget := -1 + bestSwapJ := -1 + + for gi, gRoot := range st.uniqueGroups { + grp := st.groups[gRoot] + gRoom := assignment[grp[0]] + + for room := range st.numRooms { + if room == gRoom { + continue + } + if roomCount(assignment, room)+len(grp) > roomSize { + continue + } + for _, m := range grp { + assignment[m] = room + } + if feasible(assignment) { + delta := st.score(assignment) - currentScore + if delta > bestDelta { + bestDelta = delta + bestMove = gi + bestTarget = room + bestSwapJ = -1 + } + } + for _, m := range grp { + assignment[m] = gRoom + } + } + + for gj := gi + 1; gj < len(st.uniqueGroups); gj++ { + grp2 := st.groups[st.uniqueGroups[gj]] + g2Room := assignment[grp2[0]] + if gRoom == g2Room { + continue + } + newGRoom := roomCount(assignment, gRoom) - len(grp) + len(grp2) + newG2Room := roomCount(assignment, g2Room) - len(grp2) + len(grp) + if newGRoom > roomSize || newG2Room > roomSize { + continue + } + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = gRoom + } + if feasible(assignment) { + delta := st.score(assignment) - currentScore + if delta > bestDelta { + bestDelta = delta + bestMove = gi + bestTarget = -1 + bestSwapJ = gj + } + } + for _, m := range grp { + assignment[m] = gRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + } + } + + if bestDelta <= 0 { + break + } + + grp := st.groups[st.uniqueGroups[bestMove]] + gRoom := assignment[grp[0]] + if bestSwapJ < 0 { + for _, m := range grp { + assignment[m] = bestTarget + } + } else { + grp2 := st.groups[st.uniqueGroups[bestSwapJ]] + g2Room := assignment[grp2[0]] + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = gRoom + } + } + currentScore += bestDelta + } + return currentScore + } + + perturb := func(src []int, count int) { + copy(assignment, src) + indices := rng.Perm(len(st.uniqueGroups)) + count = min(count, len(indices)) + for _, gi := range indices[:count] { + grp := st.groups[st.uniqueGroups[gi]] + oldRoom := assignment[grp[0]] + rooms := rng.Perm(st.numRooms) + for _, room := range rooms { + if room == oldRoom { + continue + } + if roomCount(assignment, room)+len(grp) > roomSize { + continue + } + for _, m := range grp { + assignment[m] = room + } + if feasible(assignment) { + break + } + for _, m := range grp { + assignment[m] = oldRoom + } + } + } + } + + copy(assignment, initialAssignment) + tracker.add(assignment, hillClimb(assignment)) + + for range params.NumRandom { + if st.randomPlacement(assignment, rng) { + tracker.add(assignment, hillClimb(assignment)) + } + } + + for range params.NumPerturb { + src := tracker.bestSolutions[rng.Intn(len(tracker.bestSolutions))] + perturb(src, params.PerturbMin+rng.Intn(params.PerturbMax-params.PerturbMin)) + tracker.add(assignment, hillClimb(assignment)) + } + + results := make([]Solution, len(tracker.bestSolutions)) + for i, sol := range tracker.bestSolutions { + results[i] = Solution{Assignment: sol, Score: tracker.bestScore} + } + return results +} + +func SolveFast(n, roomSize, pnMultiple, npCost int, constraints []Constraint, params Params, rng *rand.Rand) []Solution { + if n == 0 { + return nil + } + + st := newSolverState(n, roomSize, pnMultiple, npCost, constraints) + if st.hasHardConflict() { + return nil + } + + assignment := make([]int, n) + if !st.initialPlacement(assignment) { + for i := range n { + assignment[i] = i % st.numRooms + } + } + + initialAssignment := slices.Clone(assignment) + tracker := newTracker(assignment, st.score(assignment)) + + roomCount := func(a []int, room int) int { + c := 0 + for _, r := range a { + if r == room { + c++ + } + } + return c + } + + feasible := func(a []int) bool { + for p := range st.mustApart { + if a[p[0]] == a[p[1]] { + return false + } + } + rc := map[int]int{} + for _, room := range a { + rc[room]++ + } + for _, cnt := range rc { + if cnt > roomSize { + return false + } + } + return true + } + + perturb := func(src []int, count int) { + copy(assignment, src) + indices := rng.Perm(len(st.uniqueGroups)) + count = min(count, len(indices)) + for _, gi := range indices[:count] { + grp := st.groups[st.uniqueGroups[gi]] + oldRoom := assignment[grp[0]] + rooms := rng.Perm(st.numRooms) + for _, room := range rooms { + if room == oldRoom { + continue + } + if roomCount(assignment, room)+len(grp) > roomSize { + continue + } + for _, m := range grp { + assignment[m] = room + } + if feasible(assignment) { + break + } + for _, m := range grp { + assignment[m] = oldRoom + } + } + } + } + + copy(assignment, initialAssignment) + tracker.add(assignment, st.fastHillClimb(assignment)) + + for range params.NumRandom { + if st.randomPlacement(assignment, rng) { + tracker.add(assignment, st.fastHillClimb(assignment)) + } + } + + for range params.NumPerturb { + src := tracker.bestSolutions[rng.Intn(len(tracker.bestSolutions))] + perturb(src, params.PerturbMin+rng.Intn(params.PerturbMax-params.PerturbMin)) + tracker.add(assignment, st.fastHillClimb(assignment)) + } + + results := make([]Solution, len(tracker.bestSolutions)) + for i, sol := range tracker.bestSolutions { + results[i] = Solution{Assignment: sol, Score: tracker.bestScore} + } + return results +} + +func SolveSA(n, roomSize, pnMultiple, npCost int, constraints []Constraint, params SAParams, rng *rand.Rand) []Solution { + if n == 0 { + return nil + } + + st := newSolverState(n, roomSize, pnMultiple, npCost, constraints) + if st.hasHardConflict() { + return nil + } + + assignment := make([]int, n) + if !st.initialPlacement(assignment) { + for i := range n { + assignment[i] = i % st.numRooms + } + } + + tracker := newTracker(assignment, st.score(assignment)) + counts := st.roomCounts(assignment) + + tryMove := func(assignment []int, counts []int, rng *rand.Rand) (gi int, oldRoom int, newRoom int, swapGi int, ok bool) { + nGroups := len(st.uniqueGroups) + gi = rng.Intn(nGroups) + grp := st.groups[st.uniqueGroups[gi]] + oldRoom = assignment[grp[0]] + + if rng.Intn(3) == 0 && nGroups > 1 { + swapGi = rng.Intn(nGroups - 1) + if swapGi >= gi { + swapGi++ + } + grp2 := st.groups[st.uniqueGroups[swapGi]] + g2Room := assignment[grp2[0]] + if oldRoom == g2Room { + return 0, 0, 0, 0, false + } + newOldCount := counts[oldRoom] - len(grp) + len(grp2) + newG2Count := counts[g2Room] - len(grp2) + len(grp) + if newOldCount > roomSize || newG2Count > roomSize { + return 0, 0, 0, 0, false + } + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = oldRoom + } + if !st.feasibleForGroup(assignment, st.uniqueGroups[gi], g2Room) || + !st.feasibleForGroup(assignment, st.uniqueGroups[swapGi], oldRoom) { + for _, m := range grp { + assignment[m] = oldRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + return 0, 0, 0, 0, false + } + for _, m := range grp { + assignment[m] = oldRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + return gi, oldRoom, g2Room, swapGi, true + } + + newRoom = rng.Intn(st.numRooms - 1) + if newRoom >= oldRoom { + newRoom++ + } + if counts[newRoom]+len(grp) > roomSize { + return 0, 0, 0, 0, false + } + for _, m := range grp { + assignment[m] = newRoom + } + if !st.feasibleForGroup(assignment, st.uniqueGroups[gi], newRoom) { + for _, m := range grp { + assignment[m] = oldRoom + } + return 0, 0, 0, 0, false + } + for _, m := range grp { + assignment[m] = oldRoom + } + return gi, oldRoom, newRoom, -1, true + } + + applyMove := func(assignment []int, counts []int, gi, oldRoom, newRoom, swapGi int) { + grp := st.groups[st.uniqueGroups[gi]] + if swapGi >= 0 { + grp2 := st.groups[st.uniqueGroups[swapGi]] + g2Room := assignment[grp2[0]] + counts[oldRoom] -= len(grp) + counts[g2Room] -= len(grp2) + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = oldRoom + } + counts[g2Room] += len(grp) + counts[oldRoom] += len(grp2) + } else { + counts[oldRoom] -= len(grp) + for _, m := range grp { + assignment[m] = newRoom + } + counts[newRoom] += len(grp) + } + } + + for restart := range params.Restarts { + if restart > 0 { + if !st.randomPlacement(assignment, rng) { + continue + } + counts = st.roomCounts(assignment) + } + + currentScore := st.score(assignment) + + for step := range params.Steps { + t := params.TempHigh * math.Pow(params.TempLow/params.TempHigh, float64(step)/float64(params.Steps-1)) + + gi, oldRoom, newRoom, swapGi, ok := tryMove(assignment, counts, rng) + if !ok { + continue + } + + applyMove(assignment, counts, gi, oldRoom, newRoom, swapGi) + newScore := st.score(assignment) + delta := newScore - currentScore + + if delta >= 0 || rng.Float64() < math.Exp(float64(delta)/t) { + currentScore = newScore + tracker.add(assignment, currentScore) + } else { + if swapGi >= 0 { + grp := st.groups[st.uniqueGroups[gi]] + grp2 := st.groups[st.uniqueGroups[swapGi]] + g2Room := assignment[grp2[0]] + counts[g2Room] -= len(grp2) + counts[oldRoom] -= len(grp) + for _, m := range grp { + assignment[m] = oldRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + counts[oldRoom] += len(grp) + counts[g2Room] += len(grp2) + } else { + grp := st.groups[st.uniqueGroups[gi]] + counts[newRoom] -= len(grp) + for _, m := range grp { + assignment[m] = oldRoom + } + counts[oldRoom] += len(grp) + } + } + } + } + + results := make([]Solution, len(tracker.bestSolutions)) + for i, sol := range tracker.bestSolutions { + results[i] = Solution{Assignment: sol, Score: tracker.bestScore} + } + return results +} + +func SolveHybrid(n, roomSize, pnMultiple, npCost int, constraints []Constraint, params HybridParams, rng *rand.Rand) []Solution { + if n == 0 { + return nil + } + + st := newSolverState(n, roomSize, pnMultiple, npCost, constraints) + if st.hasHardConflict() { + return nil + } + + numRooms := st.numRooms + + roomCount := func(a []int, room int) int { + c := 0 + for _, r := range a { + if r == room { + c++ + } + } + return c + } + + feasible := func(a []int) bool { + for p := range st.mustApart { + if a[p[0]] == a[p[1]] { + return false + } + } + rc := map[int]int{} + for _, room := range a { + rc[room]++ + } + for _, cnt := range rc { + if cnt > roomSize { + return false + } + } + return true + } + + hillClimb := func(assignment []int) int { + currentScore := st.score(assignment) + for { + bestDelta := 0 + bestMove := -1 + bestTarget := -1 + bestSwapJ := -1 + + for gi, gRoot := range st.uniqueGroups { + grp := st.groups[gRoot] + gRoom := assignment[grp[0]] + + for room := range numRooms { + if room == gRoom { + continue + } + if roomCount(assignment, room)+len(grp) > roomSize { + continue + } + for _, m := range grp { + assignment[m] = room + } + if feasible(assignment) { + delta := st.score(assignment) - currentScore + if delta > bestDelta { + bestDelta = delta + bestMove = gi + bestTarget = room + bestSwapJ = -1 + } + } + for _, m := range grp { + assignment[m] = gRoom + } + } + + for gj := gi + 1; gj < len(st.uniqueGroups); gj++ { + grp2 := st.groups[st.uniqueGroups[gj]] + g2Room := assignment[grp2[0]] + if gRoom == g2Room { + continue + } + newGRoom := roomCount(assignment, gRoom) - len(grp) + len(grp2) + newG2Room := roomCount(assignment, g2Room) - len(grp2) + len(grp) + if newGRoom > roomSize || newG2Room > roomSize { + continue + } + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = gRoom + } + if feasible(assignment) { + delta := st.score(assignment) - currentScore + if delta > bestDelta { + bestDelta = delta + bestMove = gi + bestTarget = -1 + bestSwapJ = gj + } + } + for _, m := range grp { + assignment[m] = gRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + } + } + + if bestDelta <= 0 { + break + } + + grp := st.groups[st.uniqueGroups[bestMove]] + gRoom := assignment[grp[0]] + if bestSwapJ < 0 { + for _, m := range grp { + assignment[m] = bestTarget + } + } else { + grp2 := st.groups[st.uniqueGroups[bestSwapJ]] + g2Room := assignment[grp2[0]] + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = gRoom + } + } + currentScore += bestDelta + } + return currentScore + } + + assignment := make([]int, n) + if !st.initialPlacement(assignment) { + for i := range n { + assignment[i] = i % numRooms + } + } + + tracker := newTracker(assignment, st.score(assignment)) + tracker.add(assignment, hillClimb(assignment)) + + counts := make([]int, numRooms) + + for restart := range params.SARestarts { + if restart == 0 { + copy(assignment, tracker.bestSolutions[0]) + } else { + if !st.randomPlacement(assignment, rng) { + continue + } + } + for i := range counts { + counts[i] = 0 + } + for _, room := range assignment { + counts[room]++ + } + + currentScore := st.score(assignment) + + for step := range params.SASteps { + t := params.TempHigh * math.Pow(params.TempLow/params.TempHigh, float64(step)/float64(params.SASteps-1)) + + nGroups := len(st.uniqueGroups) + gi := rng.Intn(nGroups) + grp := st.groups[st.uniqueGroups[gi]] + oldRoom := assignment[grp[0]] + + var newRoom int + swapGi := -1 + moved := false + + if rng.Intn(3) == 0 && nGroups > 1 { + swapGi = rng.Intn(nGroups - 1) + if swapGi >= gi { + swapGi++ + } + grp2 := st.groups[st.uniqueGroups[swapGi]] + g2Room := assignment[grp2[0]] + if oldRoom != g2Room { + newOld := counts[oldRoom] - len(grp) + len(grp2) + newG2 := counts[g2Room] - len(grp2) + len(grp) + if newOld <= roomSize && newG2 <= roomSize { + for _, m := range grp { + assignment[m] = g2Room + } + for _, m := range grp2 { + assignment[m] = oldRoom + } + if st.feasibleForGroup(assignment, st.uniqueGroups[gi], g2Room) && + st.feasibleForGroup(assignment, st.uniqueGroups[swapGi], oldRoom) { + newRoom = g2Room + moved = true + counts[oldRoom] = newOld + counts[g2Room] = newG2 + } else { + for _, m := range grp { + assignment[m] = oldRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + } + } + } + } + + if !moved { + swapGi = -1 + newRoom = rng.Intn(numRooms - 1) + if newRoom >= oldRoom { + newRoom++ + } + if counts[newRoom]+len(grp) <= roomSize { + for _, m := range grp { + assignment[m] = newRoom + } + if st.feasibleForGroup(assignment, st.uniqueGroups[gi], newRoom) { + moved = true + counts[oldRoom] -= len(grp) + counts[newRoom] += len(grp) + } else { + for _, m := range grp { + assignment[m] = oldRoom + } + } + } + } + + if !moved { + continue + } + + newScore := st.score(assignment) + delta := newScore - currentScore + + if delta >= 0 || rng.Float64() < math.Exp(float64(delta)/t) { + currentScore = newScore + } else { + if swapGi >= 0 { + grp2 := st.groups[st.uniqueGroups[swapGi]] + g2Room := assignment[grp2[0]] + for _, m := range grp { + assignment[m] = oldRoom + } + for _, m := range grp2 { + assignment[m] = g2Room + } + counts[oldRoom] = counts[oldRoom] - len(grp2) + len(grp) + counts[g2Room] = counts[g2Room] - len(grp) + len(grp2) + } else { + for _, m := range grp { + assignment[m] = oldRoom + } + counts[newRoom] -= len(grp) + counts[oldRoom] += len(grp) + } + } + } + + tracker.add(assignment, hillClimb(assignment)) + } + + results := make([]Solution, len(tracker.bestSolutions)) + for i, sol := range tracker.bestSolutions { + results[i] = Solution{Assignment: sol, Score: tracker.bestScore} + } + return results +}