diff --git a/ortools/util/BUILD.bazel b/ortools/util/BUILD.bazel
index 4b964ab7f5..6105c72214 100644
--- a/ortools/util/BUILD.bazel
+++ b/ortools/util/BUILD.bazel
@@ -358,6 +358,7 @@ cc_library(
     name = "vector_or_function",
     hdrs = ["vector_or_function.h"],
     deps = [
+        ":flat_matrix",
         "//ortools/base",
     ],
 )
@@ -405,6 +406,7 @@ cc_library(
 cc_library(
     name = "random_engine",
     hdrs = ["random_engine.h"],
+    deps = [],
 )
 
 cc_library(
diff --git a/ortools/util/vector_or_function.h b/ortools/util/vector_or_function.h
index 8a02380a49..103cd0c0ef 100644
--- a/ortools/util/vector_or_function.h
+++ b/ortools/util/vector_or_function.h
@@ -18,6 +18,7 @@
 #include <vector>
 
 #include "ortools/base/logging.h"
+#include "ortools/util/flat_matrix.h"
 
 namespace operations_research {
 
@@ -90,6 +91,20 @@ class MatrixOrFunction<ScalarType, std::vector<std::vector<ScalarType>>,
   std::vector<std::vector<ScalarType>> matrix_;
 };
 
+// Specialization for FlatMatrix<>, which is faster than vector<vector<>>.
+template <typename ScalarType, bool square>
+class MatrixOrFunction<ScalarType, FlatMatrix<ScalarType>, square> {
+ public:
+  explicit MatrixOrFunction(FlatMatrix<ScalarType> matrix)
+      : matrix_(std::move(matrix)) {}
+  void Reset(FlatMatrix<ScalarType> matrix) { matrix_ = std::move(matrix); }
+  ScalarType operator()(int i, int j) const { return matrix_[i][j]; }
+  bool Check() const { return true; }
+
+ private:
+  FlatMatrix<ScalarType> matrix_;
+};
+
 }  // namespace operations_research
 
 #endif  // OR_TOOLS_UTIL_VECTOR_OR_FUNCTION_H_