improve cp-sat python types
This commit is contained in:
@@ -726,7 +726,7 @@ def create_data_model() -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
return capacity_df, tasks_df
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
"""Create the model and solves it."""
|
||||
capacity_df, tasks_df = create_data_model()
|
||||
|
||||
@@ -834,7 +834,7 @@ def rank_tasks(
|
||||
starts: list[cp_model.IntVar],
|
||||
presences: list[cp_model.IntVar],
|
||||
ranks: list[cp_model.IntVar],
|
||||
):
|
||||
) -> None:
|
||||
"""This method adds constraints and variables to links tasks and ranks.
|
||||
|
||||
This method assumes that all starts are disjoint, meaning that all tasks have
|
||||
@@ -852,7 +852,7 @@ def rank_tasks(
|
||||
all_tasks = range(num_tasks)
|
||||
|
||||
# Creates precedence variables between pairs of intervals.
|
||||
precedences = {}
|
||||
precedences: dict[tuple[int, int], cp_model.IntVar] = {}
|
||||
for i in all_tasks:
|
||||
for j in all_tasks:
|
||||
if i == j:
|
||||
@@ -865,7 +865,10 @@ def rank_tasks(
|
||||
# Treats optional intervals.
|
||||
for i in range(num_tasks - 1):
|
||||
for j in range(i + 1, num_tasks):
|
||||
tmp_array = [precedences[(i, j)], precedences[(j, i)]]
|
||||
tmp_array: list[cp_model.LiteralT] = [
|
||||
precedences[(i, j)],
|
||||
precedences[(j, i)],
|
||||
]
|
||||
if not cp_model.object_is_a_true_literal(presences[i]):
|
||||
tmp_array.append(presences[i].negated())
|
||||
# Makes sure that if i is not performed, all precedences are false.
|
||||
@@ -898,7 +901,7 @@ def rank_tasks(
|
||||
model.add(ranks[i] == sum(precedences[(j, i)] for j in all_tasks) - 1)
|
||||
|
||||
|
||||
def ranking_sample_sat():
|
||||
def ranking_sample_sat() -> None:
|
||||
"""Ranks tasks in a NoOverlap constraint."""
|
||||
|
||||
model = cp_model.CpModel()
|
||||
|
||||
@@ -101,7 +101,7 @@ parallelism. Therefore, the number of workers must be set to 1.
|
||||
from ortools.sat.python import cp_model
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
"""Showcases assumptions."""
|
||||
# Creates the model.
|
||||
model = cp_model.CpModel()
|
||||
|
||||
@@ -56,6 +56,7 @@ from typing import (
|
||||
Dict,
|
||||
Iterable,
|
||||
List,
|
||||
NoReturn,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
@@ -80,48 +81,30 @@ Domain = sorted_interval_list.Domain
|
||||
# usual arithmetic operators + - * / and with constant numbers, which makes the
|
||||
# python API very intuitive. See../ samples/*.py for examples.
|
||||
|
||||
INT_MIN = -9223372036854775808 # hardcoded to be platform independent.
|
||||
INT_MAX = 9223372036854775807
|
||||
INT32_MAX = 2147483647
|
||||
INT32_MIN = -2147483648
|
||||
INT_MIN = -(2**63) # hardcoded to be platform independent.
|
||||
INT_MAX = 2**63 - 1
|
||||
INT32_MIN = -(2**31)
|
||||
INT32_MAX = 2**31 - 1
|
||||
|
||||
# CpSolver status (exported to avoid importing cp_model_cp2).
|
||||
UNKNOWN: cp_model_pb2.CpSolverStatus = cp_model_pb2.UNKNOWN
|
||||
MODEL_INVALID: cp_model_pb2.CpSolverStatus = cp_model_pb2.MODEL_INVALID
|
||||
FEASIBLE: cp_model_pb2.CpSolverStatus = cp_model_pb2.FEASIBLE
|
||||
INFEASIBLE: cp_model_pb2.CpSolverStatus = cp_model_pb2.INFEASIBLE
|
||||
OPTIMAL: cp_model_pb2.CpSolverStatus = cp_model_pb2.OPTIMAL
|
||||
UNKNOWN = cp_model_pb2.UNKNOWN
|
||||
MODEL_INVALID = cp_model_pb2.MODEL_INVALID
|
||||
FEASIBLE = cp_model_pb2.FEASIBLE
|
||||
INFEASIBLE = cp_model_pb2.INFEASIBLE
|
||||
OPTIMAL = cp_model_pb2.OPTIMAL
|
||||
|
||||
# Variable selection strategy
|
||||
CHOOSE_FIRST: cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy = (
|
||||
cp_model_pb2.DecisionStrategyProto.CHOOSE_FIRST
|
||||
)
|
||||
CHOOSE_LOWEST_MIN: (
|
||||
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
|
||||
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_LOWEST_MIN
|
||||
CHOOSE_HIGHEST_MAX: (
|
||||
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
|
||||
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_HIGHEST_MAX
|
||||
CHOOSE_MIN_DOMAIN_SIZE: (
|
||||
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
|
||||
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_MIN_DOMAIN_SIZE
|
||||
CHOOSE_MAX_DOMAIN_SIZE: (
|
||||
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
|
||||
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_MAX_DOMAIN_SIZE
|
||||
CHOOSE_FIRST = cp_model_pb2.DecisionStrategyProto.CHOOSE_FIRST
|
||||
CHOOSE_LOWEST_MIN = cp_model_pb2.DecisionStrategyProto.CHOOSE_LOWEST_MIN
|
||||
CHOOSE_HIGHEST_MAX = cp_model_pb2.DecisionStrategyProto.CHOOSE_HIGHEST_MAX
|
||||
CHOOSE_MIN_DOMAIN_SIZE = cp_model_pb2.DecisionStrategyProto.CHOOSE_MIN_DOMAIN_SIZE
|
||||
CHOOSE_MAX_DOMAIN_SIZE = cp_model_pb2.DecisionStrategyProto.CHOOSE_MAX_DOMAIN_SIZE
|
||||
|
||||
# Domain reduction strategy
|
||||
SELECT_MIN_VALUE: cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy = (
|
||||
cp_model_pb2.DecisionStrategyProto.SELECT_MIN_VALUE
|
||||
)
|
||||
SELECT_MAX_VALUE: cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy = (
|
||||
cp_model_pb2.DecisionStrategyProto.SELECT_MAX_VALUE
|
||||
)
|
||||
SELECT_LOWER_HALF: (
|
||||
cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy
|
||||
) = cp_model_pb2.DecisionStrategyProto.SELECT_LOWER_HALF
|
||||
SELECT_UPPER_HALF: (
|
||||
cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy
|
||||
) = cp_model_pb2.DecisionStrategyProto.SELECT_UPPER_HALF
|
||||
SELECT_MIN_VALUE = cp_model_pb2.DecisionStrategyProto.SELECT_MIN_VALUE
|
||||
SELECT_MAX_VALUE = cp_model_pb2.DecisionStrategyProto.SELECT_MAX_VALUE
|
||||
SELECT_LOWER_HALF = cp_model_pb2.DecisionStrategyProto.SELECT_LOWER_HALF
|
||||
SELECT_UPPER_HALF = cp_model_pb2.DecisionStrategyProto.SELECT_UPPER_HALF
|
||||
|
||||
# Search branching
|
||||
AUTOMATIC_SEARCH = sat_parameters_pb2.SatParameters.AUTOMATIC_SEARCH
|
||||
@@ -144,6 +127,7 @@ BoolVarT = Union["IntVar", "_NotBooleanVariable"]
|
||||
VariableT = Union["IntVar", IntegralT]
|
||||
LinearExprT = Union["LinearExpr", "IntVar", IntegralT]
|
||||
ObjLinearExprT = Union["LinearExpr", "IntVar", NumberT]
|
||||
BoundedLinearExprT = Union["BoundedLinearExpression", bool]
|
||||
ArcT = Tuple[IntegralT, IntegralT, LiteralT]
|
||||
_IndexOrSeries = Union[pd.Index, pd.Series]
|
||||
|
||||
@@ -405,34 +389,34 @@ class LinearExpr:
|
||||
break
|
||||
return coeffs, constant, is_integer
|
||||
|
||||
def __hash__(self):
|
||||
def __hash__(self) -> int:
|
||||
return object.__hash__(self)
|
||||
|
||||
def __abs__(self):
|
||||
def __abs__(self) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling abs() on a linear expression is not supported, "
|
||||
"please use CpModel.add_abs_equality"
|
||||
)
|
||||
|
||||
def __add__(self, arg):
|
||||
def __add__(self, arg) -> LinearExprT:
|
||||
if cmh.is_zero(arg):
|
||||
return self
|
||||
return _Sum(self, arg)
|
||||
|
||||
def __radd__(self, arg):
|
||||
def __radd__(self, arg) -> LinearExprT:
|
||||
if cmh.is_zero(arg):
|
||||
return self
|
||||
return _Sum(self, arg)
|
||||
|
||||
def __sub__(self, arg):
|
||||
def __sub__(self, arg) -> LinearExprT:
|
||||
if cmh.is_zero(arg):
|
||||
return self
|
||||
return _Sum(self, -arg)
|
||||
|
||||
def __rsub__(self, arg):
|
||||
def __rsub__(self, arg) -> LinearExprT:
|
||||
return _Sum(-self, arg)
|
||||
|
||||
def __mul__(self, arg):
|
||||
def __mul__(self, arg) -> LinearExprT:
|
||||
arg = cmh.assert_is_a_number(arg)
|
||||
if cmh.is_one(arg):
|
||||
return self
|
||||
@@ -440,7 +424,7 @@ class LinearExpr:
|
||||
return 0
|
||||
return _ProductCst(self, arg)
|
||||
|
||||
def __rmul__(self, arg):
|
||||
def __rmul__(self, arg) -> LinearExprT:
|
||||
arg = cmh.assert_is_a_number(arg)
|
||||
if cmh.is_one(arg):
|
||||
return self
|
||||
@@ -448,67 +432,67 @@ class LinearExpr:
|
||||
return 0
|
||||
return _ProductCst(self, arg)
|
||||
|
||||
def __div__(self, _):
|
||||
def __div__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling / on a linear expression is not supported, "
|
||||
"please use CpModel.add_division_equality"
|
||||
)
|
||||
|
||||
def __truediv__(self, _):
|
||||
def __truediv__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling // on a linear expression is not supported, "
|
||||
"please use CpModel.add_division_equality"
|
||||
)
|
||||
|
||||
def __mod__(self, _):
|
||||
def __mod__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling %% on a linear expression is not supported, "
|
||||
"please use CpModel.add_modulo_equality"
|
||||
)
|
||||
|
||||
def __pow__(self, _):
|
||||
def __pow__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling ** on a linear expression is not supported, "
|
||||
"please use CpModel.add_multiplication_equality"
|
||||
)
|
||||
|
||||
def __lshift__(self, _):
|
||||
def __lshift__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling left shift on a linear expression is not supported"
|
||||
)
|
||||
|
||||
def __rshift__(self, _):
|
||||
def __rshift__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling right shift on a linear expression is not supported"
|
||||
)
|
||||
|
||||
def __and__(self, _):
|
||||
def __and__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling and on a linear expression is not supported, "
|
||||
"please use CpModel.add_bool_and"
|
||||
)
|
||||
|
||||
def __or__(self, _):
|
||||
def __or__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling or on a linear expression is not supported, "
|
||||
"please use CpModel.add_bool_or"
|
||||
)
|
||||
|
||||
def __xor__(self, _):
|
||||
def __xor__(self, _) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"calling xor on a linear expression is not supported, "
|
||||
"please use CpModel.add_bool_xor"
|
||||
)
|
||||
|
||||
def __neg__(self):
|
||||
def __neg__(self) -> LinearExprT:
|
||||
return _ProductCst(self, -1)
|
||||
|
||||
def __bool__(self):
|
||||
def __bool__(self) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"Evaluating a LinearExpr instance as a Boolean is not implemented."
|
||||
)
|
||||
|
||||
def __eq__(self, arg):
|
||||
def __eq__(self, arg) -> BoundedLinearExprT:
|
||||
if arg is None:
|
||||
return False
|
||||
if cmh.is_integral(arg):
|
||||
@@ -517,21 +501,21 @@ class LinearExpr:
|
||||
else:
|
||||
return BoundedLinearExpression(self - arg, [0, 0])
|
||||
|
||||
def __ge__(self, arg):
|
||||
def __ge__(self, arg) -> BoundedLinearExprT:
|
||||
if cmh.is_integral(arg):
|
||||
arg = cmh.assert_is_int64(arg)
|
||||
return BoundedLinearExpression(self, [arg, INT_MAX])
|
||||
else:
|
||||
return BoundedLinearExpression(self - arg, [0, INT_MAX])
|
||||
|
||||
def __le__(self, arg):
|
||||
def __le__(self, arg) -> BoundedLinearExprT:
|
||||
if cmh.is_integral(arg):
|
||||
arg = cmh.assert_is_int64(arg)
|
||||
return BoundedLinearExpression(self, [INT_MIN, arg])
|
||||
else:
|
||||
return BoundedLinearExpression(self - arg, [INT_MIN, 0])
|
||||
|
||||
def __lt__(self, arg):
|
||||
def __lt__(self, arg) -> BoundedLinearExprT:
|
||||
if cmh.is_integral(arg):
|
||||
arg = cmh.assert_is_int64(arg)
|
||||
if arg == INT_MIN:
|
||||
@@ -540,7 +524,7 @@ class LinearExpr:
|
||||
else:
|
||||
return BoundedLinearExpression(self - arg, [INT_MIN, -1])
|
||||
|
||||
def __gt__(self, arg):
|
||||
def __gt__(self, arg) -> BoundedLinearExprT:
|
||||
if cmh.is_integral(arg):
|
||||
arg = cmh.assert_is_int64(arg)
|
||||
if arg == INT_MAX:
|
||||
@@ -549,7 +533,7 @@ class LinearExpr:
|
||||
else:
|
||||
return BoundedLinearExpression(self - arg, [1, INT_MAX])
|
||||
|
||||
def __ne__(self, arg):
|
||||
def __ne__(self, arg) -> BoundedLinearExprT:
|
||||
if arg is None:
|
||||
return True
|
||||
if cmh.is_integral(arg):
|
||||
@@ -904,7 +888,7 @@ class _NotBooleanVariable(LinearExpr):
|
||||
def name(self) -> str:
|
||||
return "not(%s)" % str(self.__boolvar)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
def __bool__(self) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"Evaluating a literal as a Boolean value is not implemented."
|
||||
)
|
||||
@@ -964,8 +948,9 @@ class BoundedLinearExpression:
|
||||
return self.__bounds
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
if isinstance(self.__expr, LinearExpr):
|
||||
coeffs_map, constant = self.__expr.get_integer_var_value_map()
|
||||
expr = self.__expr
|
||||
if isinstance(expr, LinearExpr):
|
||||
coeffs_map, constant = expr.get_integer_var_value_map()
|
||||
all_coeffs = set(coeffs_map.values())
|
||||
same_var = set([0])
|
||||
eq_bounds = [0, 0]
|
||||
@@ -3181,7 +3166,7 @@ class CpSolver:
|
||||
"""Returns the indices of the infeasible assumptions."""
|
||||
return self._solution.sufficient_assumptions_for_infeasibility
|
||||
|
||||
def status_name(self, status: Optional[cp_model_pb2.CpSolverStatus] = None) -> str:
|
||||
def status_name(self, status: Optional[Any] = None) -> str:
|
||||
"""Returns the name of the status returned by solve()."""
|
||||
if status is None:
|
||||
status = self._solution.status
|
||||
@@ -3245,7 +3230,7 @@ class CpSolver:
|
||||
def SolutionInfo(self) -> str:
|
||||
return self.solution_info()
|
||||
|
||||
def StatusName(self, status: Optional[cp_model_pb2.CpSolverStatus] = None) -> str:
|
||||
def StatusName(self, status: Optional[Any] = None) -> str:
|
||||
return self.status_name(status)
|
||||
|
||||
def StopSearch(self) -> None:
|
||||
|
||||
@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# Data
|
||||
# [START data]
|
||||
costs = [
|
||||
|
||||
@@ -24,7 +24,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# Data
|
||||
# [START data_model]
|
||||
data_str = """
|
||||
|
||||
@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# Data
|
||||
# [START data]
|
||||
costs = [
|
||||
|
||||
@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# Data
|
||||
# [START data]
|
||||
costs = [
|
||||
|
||||
@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
"""Showcases assumptions."""
|
||||
# Creates the model.
|
||||
# [START model]
|
||||
|
||||
@@ -61,7 +61,7 @@ def create_data_model() -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
# [END data_model]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# [START data]
|
||||
items, bins = create_data_model()
|
||||
# [END data]
|
||||
|
||||
@@ -46,7 +46,7 @@ class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
|
||||
# [END solution_printer]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
"""solve the CP+IS+FUN==TRUE cryptarithm."""
|
||||
# Constraint programming engine
|
||||
# [START model]
|
||||
|
||||
@@ -20,7 +20,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
"""Minimal jobshop problem."""
|
||||
# Data.
|
||||
# [START data]
|
||||
|
||||
@@ -19,18 +19,18 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# [START data]
|
||||
data = {}
|
||||
data["weights"] = [48, 30, 42, 36, 36, 48, 42, 42, 36, 24, 30, 30, 42, 36, 36]
|
||||
data["values"] = [10, 30, 25, 50, 35, 30, 15, 40, 30, 35, 45, 10, 20, 30, 25]
|
||||
assert len(data["weights"]) == len(data["values"])
|
||||
data["num_items"] = len(data["weights"])
|
||||
data["all_items"] = range(data["num_items"])
|
||||
num_items = len(data["weights"])
|
||||
all_items = range(num_items)
|
||||
|
||||
data["bin_capacities"] = [100, 100, 100, 100, 100]
|
||||
data["num_bins"] = len(data["bin_capacities"])
|
||||
data["all_bins"] = range(data["num_bins"])
|
||||
num_bins = len(data["bin_capacities"])
|
||||
all_bins = range(num_bins)
|
||||
# [END data]
|
||||
|
||||
# [START model]
|
||||
@@ -41,21 +41,21 @@ def main():
|
||||
# [START variables]
|
||||
# x[i, b] = 1 if item i is packed in bin b.
|
||||
x = {}
|
||||
for i in data["all_items"]:
|
||||
for b in data["all_bins"]:
|
||||
for i in all_items:
|
||||
for b in all_bins:
|
||||
x[i, b] = model.new_bool_var(f"x_{i}_{b}")
|
||||
# [END variables]
|
||||
|
||||
# Constraints.
|
||||
# [START constraints]
|
||||
# Each item is assigned to at most one bin.
|
||||
for i in data["all_items"]:
|
||||
model.add_at_most_one(x[i, b] for b in data["all_bins"])
|
||||
for i in all_items:
|
||||
model.add_at_most_one(x[i, b] for b in all_bins)
|
||||
|
||||
# The amount packed in each bin cannot exceed its capacity.
|
||||
for b in data["all_bins"]:
|
||||
for b in all_bins:
|
||||
model.add(
|
||||
sum(x[i, b] * data["weights"][i] for i in data["all_items"])
|
||||
sum(x[i, b] * data["weights"][i] for i in all_items)
|
||||
<= data["bin_capacities"][b]
|
||||
)
|
||||
# [END constraints]
|
||||
@@ -64,8 +64,8 @@ def main():
|
||||
# [START objective]
|
||||
# maximize total value of packed items.
|
||||
objective = []
|
||||
for i in data["all_items"]:
|
||||
for b in data["all_bins"]:
|
||||
for i in all_items:
|
||||
for b in all_bins:
|
||||
objective.append(cp_model.LinearExpr.term(x[i, b], data["values"][i]))
|
||||
model.maximize(cp_model.LinearExpr.sum(objective))
|
||||
# [END objective]
|
||||
@@ -79,11 +79,11 @@ def main():
|
||||
if status == cp_model.OPTIMAL:
|
||||
print(f"Total packed value: {solver.objective_value}")
|
||||
total_weight = 0
|
||||
for b in data["all_bins"]:
|
||||
for b in all_bins:
|
||||
print(f"Bin {b}")
|
||||
bin_weight = 0
|
||||
bin_value = 0
|
||||
for i in data["all_items"]:
|
||||
for i in all_items:
|
||||
if solver.value(x[i, b]) > 0:
|
||||
print(
|
||||
f"Item:{i} weight:{data['weights'][i]} value:{data['values'][i]}"
|
||||
|
||||
@@ -57,7 +57,7 @@ class NQueenSolutionPrinter(cp_model.CpSolverSolutionCallback):
|
||||
# [END solution_printer]
|
||||
|
||||
|
||||
def main(board_size):
|
||||
def main(board_size: int) -> None:
|
||||
# Creates the solver.
|
||||
# [START model]
|
||||
model = cp_model.CpModel()
|
||||
|
||||
@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# Data.
|
||||
# [START data]
|
||||
num_nurses = 4
|
||||
|
||||
@@ -20,9 +20,9 @@ from ortools.sat.python import cp_model
|
||||
def rank_tasks(
|
||||
model: cp_model.CpModel,
|
||||
starts: list[cp_model.IntVar],
|
||||
presences: list[cp_model.IntVar],
|
||||
presences: list[cp_model.BoolVarT],
|
||||
ranks: list[cp_model.IntVar],
|
||||
):
|
||||
) -> None:
|
||||
"""This method adds constraints and variables to links tasks and ranks.
|
||||
|
||||
This method assumes that all starts are disjoint, meaning that all tasks have
|
||||
@@ -32,7 +32,7 @@ def rank_tasks(
|
||||
Args:
|
||||
model: The CpModel to add the constraints to.
|
||||
starts: The array of starts variables of all tasks.
|
||||
presences: The array of presence variables of all tasks.
|
||||
presences: The array of presence variables or constants of all tasks.
|
||||
ranks: The array of rank variables of all tasks.
|
||||
"""
|
||||
|
||||
@@ -40,7 +40,7 @@ def rank_tasks(
|
||||
all_tasks = range(num_tasks)
|
||||
|
||||
# Creates precedence variables between pairs of intervals.
|
||||
precedences = {}
|
||||
precedences: dict[tuple[int, int], cp_model.BoolVarT] = {}
|
||||
for i in all_tasks:
|
||||
for j in all_tasks:
|
||||
if i == j:
|
||||
@@ -53,7 +53,7 @@ def rank_tasks(
|
||||
# Treats optional intervals.
|
||||
for i in range(num_tasks - 1):
|
||||
for j in range(i + 1, num_tasks):
|
||||
tmp_array: list[cp_model.LiteralT] = [
|
||||
tmp_array: list[cp_model.BoolVarT] = [
|
||||
precedences[(i, j)],
|
||||
precedences[(j, i)],
|
||||
]
|
||||
@@ -89,7 +89,7 @@ def rank_tasks(
|
||||
model.add(ranks[i] == sum(precedences[(j, i)] for j in all_tasks) - 1)
|
||||
|
||||
|
||||
def ranking_sample_sat():
|
||||
def ranking_sample_sat() -> None:
|
||||
"""Ranks tasks in a NoOverlap constraint."""
|
||||
|
||||
model = cp_model.CpModel()
|
||||
@@ -100,7 +100,7 @@ def ranking_sample_sat():
|
||||
starts = []
|
||||
ends = []
|
||||
intervals = []
|
||||
presences = []
|
||||
presences: list[cp_model.BoolVarT] = []
|
||||
ranks = []
|
||||
|
||||
# Creates intervals, half of them are optional.
|
||||
@@ -110,7 +110,7 @@ def ranking_sample_sat():
|
||||
end = model.new_int_var(0, horizon, f"end[{t}]")
|
||||
if t < num_tasks // 2:
|
||||
interval = model.new_interval_var(start, duration, end, f"interval[{t}]")
|
||||
presence = True
|
||||
presence = model.new_constant(1)
|
||||
else:
|
||||
presence = model.new_bool_var(f"presence[{t}]")
|
||||
interval = model.new_optional_interval_var(
|
||||
|
||||
@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
|
||||
# [END import]
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
# This program tries to find an optimal assignment of nurses to shifts
|
||||
# (3 shifts per day, for 7 days), subject to some constraints (see below).
|
||||
# Each nurse can request to be assigned to specific shifts.
|
||||
|
||||
Reference in New Issue
Block a user