117 lines
4.0 KiB
Python
117 lines
4.0 KiB
Python
# Copyright 2018 Gergo Rozner
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http:#www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Solves the stable marriage problem in CP-SAT."""
|
|
|
|
from ortools.sat.python import cp_model
|
|
|
|
|
|
class SolutionPrinter(cp_model.CpSolverSolutionCallback):
|
|
|
|
def __init__(self, wife, husband):
|
|
cp_model.CpSolverSolutionCallback.__init__(self)
|
|
self.__wife = wife
|
|
self.__husband = husband
|
|
self.__solution_count = 0
|
|
self.__n = len(wife)
|
|
|
|
def OnSolutionCallback(self):
|
|
self.__solution_count += 1
|
|
print("Solution %i" % self.__solution_count)
|
|
print("wife : ", [self.Value(self.__wife[i]) for i in range(self.__n)])
|
|
print("husband: ", [self.Value(self.__husband[i]) for i in range(self.__n)])
|
|
print()
|
|
|
|
def SolutionCount(self):
|
|
return self.__solution_count
|
|
|
|
|
|
def main(ranks, pair_num):
|
|
mrank = ranks["rankMen"]
|
|
wrank = ranks["rankWomen"]
|
|
|
|
n = pair_num
|
|
|
|
model = cp_model.CpModel()
|
|
|
|
wife = [model.NewIntVar(0, n - 1, "wife[%i]" % i) for i in range(n)]
|
|
husband = [model.NewIntVar(0, n - 1, "husband[%i]" % i) for i in range(n)]
|
|
|
|
for m in range(n):
|
|
model.AddElement(wife[m], husband, m)
|
|
|
|
for w in range(n):
|
|
model.AddElement(husband[w], wife, w)
|
|
|
|
#mrank[w][m] < mrank[w][husband[w]] => wrank[m][wife[m]] < wrank[m][w]
|
|
for w in range(n):
|
|
for m in range(n):
|
|
husband_rank = model.NewIntVar(1, n, "")
|
|
model.AddElement(husband[w], mrank[w], husband_rank)
|
|
|
|
wife_rank = model.NewIntVar(1, n, "")
|
|
model.AddElement(wife[m], wrank[m], wife_rank)
|
|
|
|
husband_dominated = model.NewBoolVar("")
|
|
model.Add(mrank[w][m] < husband_rank).OnlyEnforceIf(husband_dominated)
|
|
model.Add(mrank[w][m] >= husband_rank).OnlyEnforceIf(
|
|
husband_dominated.Not())
|
|
|
|
wife_dominates = model.NewBoolVar("")
|
|
model.Add(wife_rank < wrank[m][w]).OnlyEnforceIf(wife_dominates)
|
|
model.Add(wife_rank >= wrank[m][w]).OnlyEnforceIf(wife_dominates.Not())
|
|
|
|
model.AddImplication(husband_dominated, wife_dominates)
|
|
|
|
#wrank[m][w] < wrank[m][wife[m]] => mrank[w][husband[w]] < mrank[w][m]
|
|
for m in range(n):
|
|
for w in range(n):
|
|
wife_rank = model.NewIntVar(1, n, "")
|
|
model.AddElement(wife[m], wrank[m], wife_rank)
|
|
|
|
husband_rank = model.NewIntVar(1, n, "")
|
|
model.AddElement(husband[w], mrank[w], husband_rank)
|
|
|
|
wife_dominated = model.NewBoolVar("")
|
|
model.Add(wrank[m][w] < wife_rank).OnlyEnforceIf(wife_dominated)
|
|
model.Add(wrank[m][w] >= wife_rank).OnlyEnforceIf(wife_dominated.Not())
|
|
|
|
husband_dominates = model.NewBoolVar("")
|
|
model.Add(husband_rank < mrank[w][m]).OnlyEnforceIf(husband_dominates)
|
|
model.Add(husband_rank >= mrank[w][m]).OnlyEnforceIf(
|
|
husband_dominates.Not())
|
|
|
|
model.AddImplication(wife_dominated, husband_dominates)
|
|
|
|
solver = cp_model.CpSolver()
|
|
solution_printer = SolutionPrinter(wife, husband)
|
|
solver.SearchForAllSolutions(model, solution_printer)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
rankings1 = {
|
|
"rankMen": [[1, 2, 3, 4], [4, 3, 2, 1], [1, 2, 3, 4], [3, 4, 1, 2]],
|
|
"rankWomen": [[1, 2, 3, 4], [2, 1, 3, 4], [1, 4, 3, 2], [4, 3, 1, 2]]
|
|
}
|
|
|
|
rankings2 = {
|
|
"rankMen": [[1, 5, 4, 6, 2, 3], [4, 1, 5, 2, 6, 3], [6, 4, 2, 1, 5, 3],
|
|
[1, 5, 2, 4, 3, 6], [4, 2, 1, 5, 6, 3], [2, 6, 3, 5, 1, 4]],
|
|
"rankWomen": [[1, 4, 2, 5, 6, 3], [3, 4, 6, 1, 5, 2], [1, 6, 4, 2, 3, 5],
|
|
[6, 5, 3, 4, 2, 1], [3, 1, 2, 4, 5, 6], [2, 3, 1, 6, 5, 4]]
|
|
}
|
|
|
|
problem = rankings2
|
|
couple_count = len(problem["rankMen"])
|
|
|
|
main(problem, couple_count)
|