246 lines
9.3 KiB
C++
246 lines
9.3 KiB
C++
// Copyright 2010-2025 Google LLC
|
||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||
// you may not use this file except in compliance with the License.
|
||
// You may obtain a copy of the License at
|
||
//
|
||
// http://www.apache.org/licenses/LICENSE-2.0
|
||
//
|
||
// Unless required by applicable law or agreed to in writing, software
|
||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
// See the License for the specific language governing permissions and
|
||
// limitations under the License.
|
||
|
||
#include "ortools/algorithms/n_choose_k.h"
|
||
|
||
#include <cmath>
|
||
#include <cstdint>
|
||
#include <limits>
|
||
|
||
#include "absl/numeric/int128.h"
|
||
#include "absl/random/distributions.h"
|
||
#include "absl/random/random.h"
|
||
#include "absl/status/status.h"
|
||
#include "absl/status/statusor.h"
|
||
#include "gtest/gtest.h"
|
||
#include "ortools/base/dump_vars.h"
|
||
#include "ortools/base/gmock.h"
|
||
#include "ortools/util/flat_matrix.h"
|
||
|
||
namespace operations_research {
|
||
namespace {
|
||
using ::testing::HasSubstr;
|
||
using ::testing::status::IsOkAndHolds;
|
||
using ::testing::status::StatusIs;
|
||
|
||
constexpr int64_t kint64max = std::numeric_limits<int64_t>::max();
|
||
|
||
TEST(NChooseKTest, TrivialErrorCases) {
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 100'000;
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t x = absl::LogUniform<int64_t>(random, 0, kint64max);
|
||
EXPECT_THAT(NChooseK(-1, x), StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("n is negative")));
|
||
EXPECT_THAT(NChooseK(x, -1), StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("k is negative")));
|
||
if (x != kint64max) EXPECT_THAT(NChooseK(x, x + 1), IsOkAndHolds(0));
|
||
ASSERT_FALSE(HasFailure()) << DUMP_VARS(t, x);
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, Symmetry) {
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 1'000'000;
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 0, kint64max);
|
||
const int64_t k = absl::LogUniform<int64_t>(random, 0, n);
|
||
const absl::StatusOr<int64_t> result1 = NChooseK(n, k);
|
||
const absl::StatusOr<int64_t> result2 = NChooseK(n, n - k);
|
||
if (result1.ok()) {
|
||
ASSERT_THAT(result2, IsOkAndHolds(result1.value())) << DUMP_VARS(t, n, k);
|
||
} else {
|
||
ASSERT_EQ(result2.status().code(), result1.status().code())
|
||
<< DUMP_VARS(t, n, k, result1, result2);
|
||
}
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, Invariant) {
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 1'000'000;
|
||
int num_tested_invariants = 0;
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 2, 100);
|
||
const int64_t k = absl::LogUniform<int64_t>(random, 1, n - 1);
|
||
const absl::StatusOr<int64_t> n_k = NChooseK(n, k);
|
||
const absl::StatusOr<int64_t> nm1_k = NChooseK(n - 1, k);
|
||
const absl::StatusOr<int64_t> nm1_km1 = NChooseK(n - 1, k - 1);
|
||
if (n_k.ok()) {
|
||
++num_tested_invariants;
|
||
ASSERT_OK(nm1_k);
|
||
ASSERT_OK(nm1_km1);
|
||
ASSERT_EQ(n_k.value(), nm1_k.value() + nm1_km1.value())
|
||
<< DUMP_VARS(t, n, k, n_k, nm1_k, nm1_km1);
|
||
}
|
||
}
|
||
EXPECT_GE(num_tested_invariants, kNumTests / 10);
|
||
}
|
||
|
||
TEST(NChooseKTest, ComparisonAgainstClosedFormsForK0) {
|
||
for (int64_t n : {int64_t{0}, int64_t{1}, kint64max}) {
|
||
EXPECT_THAT(NChooseK(n, 0), IsOkAndHolds(1)) << n;
|
||
}
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 1'000'000;
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 0, kint64max);
|
||
ASSERT_THAT(NChooseK(n, 0), IsOkAndHolds(1)) << DUMP_VARS(n, t);
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, ComparisonAgainstClosedFormsForK1) {
|
||
for (int64_t n : {int64_t{1}, kint64max}) {
|
||
EXPECT_THAT(NChooseK(n, 1), IsOkAndHolds(n));
|
||
}
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 1'000'000;
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 1, kint64max);
|
||
ASSERT_THAT(NChooseK(n, 1), IsOkAndHolds(n)) << DUMP_VARS(t);
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, ComparisonAgainstClosedFormsForK2) {
|
||
// 2^32 Choose 2 = 2^32 × (2^32-1) / 2 = 2^63 - 2^31 < kint64max,
|
||
// but (2^32+1) Choose 2 = 2^63 + 2^31 overflows.
|
||
constexpr int64_t max_n = int64_t{1} << 32;
|
||
for (int64_t n : {int64_t{2}, max_n}) {
|
||
const int64_t n_choose_2 =
|
||
static_cast<int64_t>(absl::uint128(n) * (n - 1) / 2);
|
||
EXPECT_THAT(NChooseK(n, 2), IsOkAndHolds(n_choose_2)) << DUMP_VARS(n);
|
||
}
|
||
EXPECT_THAT(NChooseK(max_n + 1, 2),
|
||
StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")));
|
||
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 100'000;
|
||
// Random valid results.
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 2, max_n);
|
||
const int64_t n_choose_2 =
|
||
static_cast<int64_t>(absl::uint128(n) * (n - 1) / 2);
|
||
ASSERT_THAT(NChooseK(n, 2), IsOkAndHolds(n_choose_2)) << DUMP_VARS(t, n);
|
||
}
|
||
// Random overflows.
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, max_n + 1, kint64max);
|
||
ASSERT_THAT(NChooseK(n, 2), StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")))
|
||
<< DUMP_VARS(t, n);
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, ComparisonAgainstClosedFormsForK3) {
|
||
// This is 1 + ∛6×2^21. Checked manually on Google's scientific calculator.
|
||
const int64_t max_n =
|
||
static_cast<int64_t>(1 + std::pow(6, 1.0 / 3) * std::pow(2, 21));
|
||
for (int64_t n : {int64_t{3}, max_n}) {
|
||
const int64_t n_choose_3 =
|
||
static_cast<int64_t>(absl::uint128(n) * (n - 1) * (n - 2) / 6);
|
||
EXPECT_THAT(NChooseK(n, 3), IsOkAndHolds(n_choose_3)) << DUMP_VARS(n);
|
||
}
|
||
EXPECT_THAT(NChooseK(max_n + 1, 3),
|
||
StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")));
|
||
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 100'000;
|
||
// Random valid results.
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 3, max_n);
|
||
const int64_t n_choose_3 =
|
||
static_cast<int64_t>(absl::uint128(n) * (n - 1) * (n - 2) / 6);
|
||
ASSERT_THAT(NChooseK(n, 3), IsOkAndHolds(n_choose_3)) << DUMP_VARS(t, n);
|
||
}
|
||
// Random overflows.
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, max_n + 1, kint64max);
|
||
ASSERT_THAT(NChooseK(n, 3), StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")))
|
||
<< DUMP_VARS(t, n);
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, ComparisonAgainstClosedFormsForK4) {
|
||
// This is 1.5 + ∜24 × 2^(63/4).
|
||
// Checked manually on Google's scientific calculator.
|
||
const int64_t max_n =
|
||
static_cast<int64_t>(1.5 + std::pow(24, 1.0 / 4) * std::pow(2, 63.0 / 4));
|
||
for (int64_t n : {int64_t{4}, max_n}) {
|
||
const int64_t n_choose_4 = static_cast<int64_t>(absl::uint128(n) * (n - 1) *
|
||
(n - 2) * (n - 3) / 24);
|
||
EXPECT_THAT(NChooseK(n, 4), IsOkAndHolds(n_choose_4)) << DUMP_VARS(n);
|
||
}
|
||
EXPECT_THAT(NChooseK(max_n + 1, 4),
|
||
StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")));
|
||
|
||
absl::BitGen random;
|
||
constexpr int kNumTests = 100'000;
|
||
// Random valid results.
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, 4, max_n);
|
||
const int64_t n_choose_4 = static_cast<int64_t>(absl::uint128(n) * (n - 1) *
|
||
(n - 2) * (n - 3) / 24);
|
||
ASSERT_THAT(NChooseK(n, 4), IsOkAndHolds(n_choose_4)) << DUMP_VARS(t, n);
|
||
}
|
||
// Random overflows.
|
||
for (int t = 0; t < kNumTests; ++t) {
|
||
const int64_t n = absl::LogUniform<int64_t>(random, max_n + 1, kint64max);
|
||
ASSERT_THAT(NChooseK(n, 4), StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")))
|
||
<< DUMP_VARS(t, n);
|
||
}
|
||
}
|
||
|
||
TEST(NChooseKTest, ComparisonAgainstPascalTriangleForK5OrAbove) {
|
||
// Fill the Pascal triangle. Use -1 for int64_t overflows. We go up to n =
|
||
// 17000 because (17000 Choose 5) ≈ 1.2e19 which overflows an int64_t.
|
||
constexpr int max_n = 17000;
|
||
FlatMatrix<int64_t> triangle(max_n + 1, max_n + 1);
|
||
for (int n = 0; n <= max_n; ++n) {
|
||
triangle[n][0] = 1;
|
||
triangle[n][n] = 1;
|
||
for (int i = 1; i < n; ++i) {
|
||
const int64_t a = triangle[n - 1][i - 1];
|
||
const int64_t b = triangle[n - 1][i];
|
||
if (a < 0 || b < 0 || absl::int128(a) + b > kint64max) {
|
||
triangle[n][i] = -1;
|
||
} else {
|
||
triangle[n][i] = a + b;
|
||
}
|
||
}
|
||
}
|
||
// Checking all 17000²/2 slots would be too expensive, so we check each
|
||
// "column" downwards until the first 10 overflows, and stop.
|
||
for (int k = 5; k < max_n; ++k) {
|
||
int num_overflows = 0;
|
||
for (int n = k + 5; n < max_n; ++n) {
|
||
if (num_overflows > 0) EXPECT_EQ(triangle[n][k], -1);
|
||
if (triangle[n][k] < 0) {
|
||
++num_overflows;
|
||
EXPECT_THAT(NChooseK(n, k), StatusIs(absl::StatusCode::kInvalidArgument,
|
||
HasSubstr("overflows int64")));
|
||
if (num_overflows > 10) break;
|
||
} else {
|
||
EXPECT_THAT(NChooseK(n, k), IsOkAndHolds(triangle[n][k]));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
} // namespace
|
||
} // namespace operations_research
|