From b620d7e89c8736919d0dabdb3ea4fcacfa69617a Mon Sep 17 00:00:00 2001 From: Laurent Perron Date: Fri, 18 Oct 2024 15:25:55 +0200 Subject: [PATCH] [CP-SAT} add support for element with affine expressions in C++ --- ortools/sat/cp_model.cc | 46 +++++++++++++++++++++++++++++++++-------- ortools/sat/cp_model.h | 18 ++++++++++++---- 2 files changed, 51 insertions(+), 13 deletions(-) diff --git a/ortools/sat/cp_model.cc b/ortools/sat/cp_model.cc index 392db732eaa..d73d235689d 100644 --- a/ortools/sat/cp_model.cc +++ b/ortools/sat/cp_model.cc @@ -927,24 +927,52 @@ Constraint CpModelBuilder::AddAllDifferent( } Constraint CpModelBuilder::AddVariableElement( - IntVar index, absl::Span variables, IntVar target) { + LinearExpr index, absl::Span variables, LinearExpr target) { ConstraintProto* const proto = cp_model_.add_constraints(); - proto->mutable_element()->set_index(index.index_); - proto->mutable_element()->set_target(target.index_); + *proto->mutable_element()->mutable_linear_index() = LinearExprToProto(index); + *proto->mutable_element()->mutable_linear_target() = + LinearExprToProto(target); for (const IntVar& var : variables) { - proto->mutable_element()->add_vars(var.index_); + *proto->mutable_element()->add_exprs() = LinearExprToProto(var); } return Constraint(proto); } -Constraint CpModelBuilder::AddElement(IntVar index, +Constraint CpModelBuilder::AddElement(LinearExpr index, absl::Span values, - IntVar target) { + LinearExpr target) { ConstraintProto* const proto = cp_model_.add_constraints(); - proto->mutable_element()->set_index(index.index_); - proto->mutable_element()->set_target(target.index_); + *proto->mutable_element()->mutable_linear_index() = LinearExprToProto(index); + *proto->mutable_element()->mutable_linear_target() = + LinearExprToProto(target); for (int64_t value : values) { - proto->mutable_element()->add_vars(IndexFromConstant(value)); + proto->mutable_element()->add_exprs()->set_offset(value); + } + return Constraint(proto); +} + +Constraint CpModelBuilder::AddElement(LinearExpr index, + absl::Span expressions, + LinearExpr target) { + ConstraintProto* const proto = cp_model_.add_constraints(); + *proto->mutable_element()->mutable_linear_index() = LinearExprToProto(index); + *proto->mutable_element()->mutable_linear_target() = + LinearExprToProto(target); + for (const LinearExpr& expr : expressions) { + *proto->mutable_element()->add_exprs() = LinearExprToProto(expr); + } + return Constraint(proto); +} + +Constraint CpModelBuilder::AddElement( + LinearExpr index, std::initializer_list expressions, + LinearExpr target) { + ConstraintProto* const proto = cp_model_.add_constraints(); + *proto->mutable_element()->mutable_linear_index() = LinearExprToProto(index); + *proto->mutable_element()->mutable_linear_target() = + LinearExprToProto(target); + for (const LinearExpr& expr : expressions) { + *proto->mutable_element()->add_exprs() = LinearExprToProto(expr); } return Constraint(proto); } diff --git a/ortools/sat/cp_model.h b/ortools/sat/cp_model.h index 70123343a0a..c9487504115 100644 --- a/ortools/sat/cp_model.h +++ b/ortools/sat/cp_model.h @@ -845,13 +845,23 @@ class CpModelBuilder { Constraint AddAllDifferent(std::initializer_list exprs); /// Adds the element constraint: variables[index] == target - Constraint AddVariableElement(IntVar index, + Constraint AddVariableElement(LinearExpr index, absl::Span variables, - IntVar target); + LinearExpr target); + + /// Adds the element constraint: expressions[index] == target. + Constraint AddElement(LinearExpr index, + absl::Span expressions, + LinearExpr target); + + /// Adds the element constraint: expressions[index] == target. + Constraint AddElement(LinearExpr index, + std::initializer_list expressions, + LinearExpr target); /// Adds the element constraint: values[index] == target - Constraint AddElement(IntVar index, absl::Span values, - IntVar target); + Constraint AddElement(LinearExpr index, absl::Span values, + LinearExpr target); /** * Adds a circuit constraint.