[CP-SAT} add support for element with affine expressions in C++

google · Oct 18, 2024 · 4c97fca · 4c97fca
1 parent a4f196a
commit 4c97fca
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 13 deletions.
diff --git 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<const IntVar> variables, IntVar target) {
+    LinearExpr index, absl::Span<const IntVar> 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<const int64_t> 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<const LinearExpr> 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<LinearExpr> 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
@@ -845,13 +845,23 @@ class CpModelBuilder {
   Constraint AddAllDifferent(std::initializer_list<LinearExpr> exprs);
 
   /// Adds the element constraint: variables[index] == target
-  Constraint AddVariableElement(IntVar index,
+  Constraint AddVariableElement(LinearExpr index,
                                 absl::Span<const IntVar> variables,
-                                IntVar target);
+                                LinearExpr target);
+
+  /// Adds the element constraint: expressions[index] == target.
+  Constraint AddElement(LinearExpr index,
+                        absl::Span<const LinearExpr> expressions,
+                        LinearExpr target);
+
+  /// Adds the element constraint: expressions[index] == target.
+  Constraint AddElement(LinearExpr index,
+                        std::initializer_list<LinearExpr> expressions,
+                        LinearExpr target);
 
   /// Adds the element constraint: values[index] == target
-  Constraint AddElement(IntVar index, absl::Span<const int64_t> values,
-                        IntVar target);
+  Constraint AddElement(LinearExpr index, absl::Span<const int64_t> values,
+                        LinearExpr target);
 
   /**
    * Adds a circuit constraint.