Add CountDistinct set

Author: odow

docs/src/reference/standard_form.md

@@ -100,6 +100,7 @@ Complements
 
 ```@docs
 AllDifferent
+CountDistinct
 ```
 
 ## Matrix sets

src/Test/test_basic_constraint.jl

@@ -110,6 +110,7 @@ _set(::Type{MOI.RootDetConeTriangle}) = MOI.RootDetConeTriangle(3)
 _set(::Type{MOI.RootDetConeSquare}) = MOI.RootDetConeSquare(3)
 _set(::Type{MOI.Complements}) = MOI.Complements(2)
 _set(::Type{MOI.AllDifferent}) = MOI.AllDifferent(3)
+_set(::Type{MOI.CountDistinct}) = MOI.CountDistinct(4)
 
 function _set(
     ::Type{MOI.Indicator{MOI.ACTIVATE_ON_ONE,MOI.LessThan{T}}},
@@ -274,6 +275,7 @@ for s in [
     :RootDetConeSquare,
     :Complements,
     :AllDifferent,
+    :CountDistinct,
 ]
     S = getfield(MOI, s)
     functions = if S <: MOI.AbstractScalarSet

src/Test/test_cpsat.jl

@@ -46,3 +46,44 @@ function setup_test(
     )
     return
 end
+
+"""
+    test_cpsat_CountDistinct(model::MOI.ModelLike, config::Config)
+
+Add a VectorOfVariables-in-CountDistinct constraint.
+"""
+function test_cpsat_CountDistinct(
+    model::MOI.ModelLike,
+    config::Config{T},
+) where {T}
+    @requires MOI.supports_constraint(
+        model,
+        MOI.VectorOfVariables,
+        MOI.CountDistinct,
+    )
+    @requires MOI.supports_add_constrained_variable(model, MOI.Integer)
+    @requires _supports(config, MOI.optimize!)
+    y = [MOI.add_constrained_variable(model, MOI.Integer()) for _ in 1:4]
+    x = first.(y)
+    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.CountDistinct(4))
+    MOI.optimize!(model)
+    x_val = round.(Int, MOI.get.(model, MOI.VariablePrimal(), x))
+    @test length(unique(x_val[2:end])) == x_val[1]
+    return
+end
+
+function setup_test(
+    ::typeof(test_cpsat_CountDistinct),
+    model::MOIU.MockOptimizer,
+    ::Config{T},
+) where {T}
+    MOIU.set_mock_optimize!(
+        model,
+        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
+            mock,
+            MOI.OPTIMAL,
+            (MOI.FEASIBLE_POINT, T[2, 0, 1, 0]),
+        ),
+    )
+    return
+end

src/Utilities/model.jl

@@ -788,6 +788,7 @@ const LessThanIndicatorZero{T} =
         MOI.LogDetConeTriangle,
         MOI.LogDetConeSquare,
         MOI.AllDifferent,
+        MOI.CountDistinct,
     ),
     (
         MOI.PowerCone,

src/sets.jl

@@ -1143,6 +1143,49 @@ struct AllDifferent <: AbstractVectorSet
     end
 end
 
+"""
+    CountDistinct(dimension::Int)
+
+The set ``\\{(n, x) \\in \\mathbb{R}^{1+d}\\}`` such that the number of distinct
+values in `x` is `n`.
+
+This constraint is sometimes called `nvalues`.
+
+## Example
+
+```julia
+model = Utilities.Model{Float64}()
+n = add_constrained_variable(model, MOI.Integer())
+x = [add_constrained_variable(model, MOI.Integer()) for _ in 1:3]
+add_constraint(model, vcat(n, x), CountDistinct(4))
+# if n == 1, then x[1] == x[2] == x[3]
+# if n == 2, then x[1] == x[2] != x[3] || x[1] != x[2] == x[3]
+# if n == 3, then x[1] != x[2] != x[3]
+```
+
+## Relationship to AllDifferent
+
+When the first element is `dimension - 1`, `CountDistinct` is equivalent to an
+[`AllDifferent`](@ref) constraint.
+
+```julia
+model = Utilities.Model{Float64}()
+x = [add_constrained_variable(model, MOI.Integer()) for _ in 1:3]
+add_constraint(model, vcat(3, x), CountDistinct(4))
+# equivalent to
+add_constraint(model, x, AllDifferent(3))
+```
+"""
+struct CountDistinct <: AbstractVectorSet
+    dimension::Int
+    function CountDistinct(dimension::Base.Integer)
+        if dimension < 1
+            throw(DimensionMismatch("Dimension of CountDistinct must be >= 1."))
+        end
+        return new(dimension)
+    end
+end
+
 # isbits types, nothing to copy
 function Base.copy(
     set::Union{
@@ -1178,6 +1221,7 @@ function Base.copy(
         Semicontinuous,
         Semiinteger,
         AllDifferent,
+        CountDistinct,
     },
 )
     return set

test/sets.jl

@@ -140,6 +140,7 @@ function test_sets_DimensionMismatch()
         (MOI.RootDetConeTriangle, 0),
         (MOI.RootDetConeSquare, 0),
         (MOI.AllDifferent, 0),
+        (MOI.CountDistinct, 1),
     )
         @test_throws DimensionMismatch S(min_dimension - 1)
         @test S(min_dimension) isa S