Add covered assignment solver

.github/workflows/ci.yaml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: bazel-contrib/[email protected]
       - run: bazel build //...
-      # - run: bazel test --test_output=errors //...
+      - run: bazel test --test_output=errors //...
       - uses: actions/upload-artifact@v4
         with:
           name: plugins-${{ matrix.platform }}

assignment/BUILD.bazel

@@ -0,0 +1,31 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "assignment",
+    srcs = ["assignment.cc"],
+    hdrs = ["assignment.h"],
+    deps = ["@com_google_absl//absl/strings:str_format"],
+)
+
+cc_library(
+    name = "covered_assignment",
+    srcs = ["covered_assignment.cc"],
+    hdrs = ["covered_assignment.h"],
+    deps = [
+        ":assignment",
+        "//base:logging",
+        "@com_google_absl//absl/strings:str_format",
+        "@ortools//ortools/sat:cp_model",
+    ],
+)
+
+cc_test(
+    name = "covered_assignment_test",
+    srcs = ["covered_assignment_test.cc"],
+    deps = [
+        ":assignment",
+        ":covered_assignment",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

assignment/assignment.cc

@@ -0,0 +1,30 @@
+#include "assignment/assignment.h"
+
+#include <stdexcept>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+
+namespace assignment {
+
+void Assignment::ValidateCosts() const {
+  // Validate the first dimension of the cost matrix.
+  if (costs_.size() != num_agents_) {
+    throw std::invalid_argument(
+        absl::StrFormat("The assignment cost matrix has an incorrect number of "
+                        "rows: %d vs. %d.",
+                        costs_.size(), num_agents_));
+  }
+
+  // Validate the second dimension of the cost matrix.
+  for (const auto& row : costs_) {
+    if (row.size() != num_tasks_) {
+      throw std::invalid_argument(
+          absl::StrFormat("The assignment cost matrix has an incorrect number "
+                          "of columns: %d vs. %d.",
+                          row.size(), num_tasks_));
+    }
+  }
+}
+
+}  // namespace assignment

assignment/assignment.h

@@ -0,0 +1,53 @@
+// The assignment class is an interface for a cost-based assignment problem.
+
+#pragma once
+
+#include <utility>
+#include <vector>
+
+namespace assignment {
+
+// Assignment interface.
+class Assignment {
+ public:
+  // Assignment item type.
+  using AssignmentItem = std::pair<int, int>;
+
+  // The assignment cost matrix should be a matrix of dimensions num_agents x
+  // num_tasks.
+  Assignment(int num_agents, int num_tasks,
+             std::vector<std::vector<double>> costs)
+      : num_agents_(num_agents),
+        num_tasks_(num_tasks),
+        costs_(std::move(costs)) {
+    Validate();
+    ValidateCosts();
+  }
+
+  Assignment(const Assignment&) = default;
+  Assignment& operator=(const Assignment&) = default;
+
+  virtual ~Assignment() = default;
+
+  // Assign the agents to the tasks.
+  virtual std::vector<AssignmentItem> Assign() const = 0;
+
+ protected:
+  // Validate the assignment problem.
+  virtual void Validate() const {};
+
+  // Number of agents.
+  int num_agents_ = 0;
+
+  // Number of tasks.
+  int num_tasks_ = 0;
+
+  // Assignment cost matrix of dimensions num_agents x num_tasks.
+  std::vector<std::vector<double>> costs_;
+
+ private:
+  // Validate the cost matrix.
+  void ValidateCosts() const;
+};
+
+}  // namespace assignment

assignment/covered_assignment.cc

@@ -0,0 +1,87 @@
+#include "assignment/covered_assignment.h"
+
+#include <stdexcept>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+#include "assignment/assignment.h"
+#include "base/logging.h"
+#include "ortools/sat/cp_model.h"
+#include "ortools/sat/cp_model.pb.h"
+#include "ortools/sat/cp_model_solver.h"
+
+namespace assignment {
+
+std::vector<Assignment::AssignmentItem> CoveredAssignment::Assign() const {
+  // Create the constraint programming model.
+  operations_research::sat::CpModelBuilder cp_model;
+
+  // Define the variables.
+  // x[i][j] is an array of boolean variables, such that x[i][j] is true if
+  // agent i is assigned to task j.
+  std::vector<std::vector<operations_research::sat::BoolVar>> x(
+      num_agents_, std::vector<operations_research::sat::BoolVar>(num_tasks_));
+  for (int i = 0; i < num_agents_; ++i) {
+    for (int j = 0; j < num_tasks_; ++j) {
+      x[i][j] = cp_model.NewBoolVar();
+    }
+  }
+
+  // Define the constraints.
+  // Each agent is assigned to one task.
+  for (int i = 0; i < num_agents_; ++i) {
+    cp_model.AddExactlyOne(x[i]);
+  }
+  // Each task is assigned to at least one agent.
+  for (int j = 0; j < num_tasks_; ++j) {
+    std::vector<operations_research::sat::BoolVar> tasks;
+    for (int i = 0; i < num_agents_; ++i) {
+      tasks.push_back(x[i][j]);
+    }
+    cp_model.AddAtLeastOne(tasks);
+  }
+
+  // Define the objective function.
+  operations_research::sat::DoubleLinearExpr total_cost;
+  for (int i = 0; i < num_agents_; ++i) {
+    for (int j = 0; j < num_tasks_; ++j) {
+      total_cost = total_cost.AddTerm(x[i][j], costs_[i][j]);
+    }
+  }
+  cp_model.Minimize(total_cost);
+
+  // Solve the assignment problem.
+  const operations_research::sat::CpSolverResponse response =
+      Solve(cp_model.Build());
+
+  // Check the feasibility of the solution.
+  if (response.status() ==
+      operations_research::sat::CpSolverStatus::INFEASIBLE) {
+    LOG(ERROR) << "Covered assignment problem is infeasible.";
+    return std::vector<AssignmentItem>();
+  }
+
+  // Record the assignemnts.
+  std::vector<AssignmentItem> assignments;
+  assignments.reserve(num_agents_);
+  for (int i = 0; i < num_agents_; ++i) {
+    for (int j = 0; j < num_tasks_; ++j) {
+      if (SolutionBooleanValue(response, x[i][j])) {
+        assignments.emplace_back(i, j);
+        break;
+      }
+    }
+  }
+  return assignments;
+}
+
+void CoveredAssignment::Validate() const {
+  // Check that there are at least as many agents as tasks to cover all tasks.
+  if (num_agents_ < num_tasks_) {
+    throw std::invalid_argument(
+        absl::StrFormat("There are fewer agents than tasks: %d vs. %d.",
+                        num_agents_, num_tasks_));
+  }
+}
+
+}  // namespace assignment

assignment/covered_assignment.h

@@ -0,0 +1,30 @@
+// The covered assignment assigns one task to each agent under the condition
+// that all tasks are assigned to at least one agent.
+
+#pragma once
+
+#include <vector>
+
+#include "assignment/assignment.h"
+
+namespace assignment {
+
+// Covered assignment.
+class CoveredAssignment : public Assignment {
+ public:
+  CoveredAssignment(int num_agents, int num_tasks,
+                    std::vector<std::vector<double>> costs)
+      : Assignment(num_agents, num_tasks, std::move(costs)) {}
+
+  CoveredAssignment(const CoveredAssignment&) = default;
+  CoveredAssignment& operator=(const CoveredAssignment&) = default;
+
+  // Assign the agents to the tasks.
+  std::vector<AssignmentItem> Assign() const override;
+
+ protected:
+  // Validate the assignment problem.
+  void Validate() const override;
+};
+
+}  // namespace assignment

assignment/covered_assignment_test.cc

@@ -0,0 +1,50 @@
+#include "assignment/covered_assignment.h"
+
+#include <gtest/gtest.h>
+
+#include <unordered_map>
+#include <vector>
+
+#include "assignment/assignment.h"
+
+namespace assignment {
+namespace {
+
+TEST(CoveredAssignmentTest, AssignUnique) {
+  constexpr int kNumAgents = 5;
+  constexpr int kNumTasks = 4;
+  const std::vector<std::vector<double>> costs{
+      {90, 80, 75, 70},   {35, 85, 55, 65},   {125, 95, 90, 95},
+      {45, 110, 95, 115}, {50, 100, 90, 100},
+  };
+  CoveredAssignment assignment(kNumAgents, kNumTasks, costs);
+  const auto assignments = assignment.Assign();
+  std::unordered_map<int, int> expected_assignments{
+      {0, 3}, {1, 2}, {2, 1}, {3, 0}, {4, 0}};
+  for (const auto& [agent_index, task_index] : assignments) {
+    EXPECT_EQ(expected_assignments[agent_index], task_index)
+        << "Agent " << agent_index << " was assigned to task " << task_index
+        << " but expected task " << expected_assignments[agent_index] << ".";
+  }
+}
+
+TEST(CoveredAssignmentTest, AssignMultiple) {
+  constexpr int kNumAgents = 3;
+  constexpr int kNumTasks = 2;
+  const std::vector<std::vector<double>> costs{
+      {1, 3},
+      {4, 3},
+      {2, 2},
+  };
+  CoveredAssignment assignment(kNumAgents, kNumTasks, costs);
+  const auto assignments = assignment.Assign();
+  std::unordered_map<int, int> expected_assignments{{0, 0}, {1, 1}, {2, 1}};
+  for (const auto& [agent_index, task_index] : assignments) {
+    EXPECT_EQ(expected_assignments[agent_index], task_index)
+        << "Agent " << agent_index << " was assigned to task " << task_index
+        << " but expected task " << expected_assignments[agent_index] << ".";
+  }
+}
+
+}  // namespace
+}  // namespace assignment