google-research

1
// Copyright 2023 The Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14

15
#include "fairness_constraint.h"
16

17
#include <algorithm>
18
#include <cassert>
19
#include <map>
20
#include <memory>
21
#include <set>
22
#include <utility>
23
#include <vector>
24

25
#include "matroid.h"
26
#include "partition_matroid.h"
27

28
FairnessConstraint::FairnessConstraint(
29
    const std::map<int, int>& colors_map,
30
    const std::vector<std::pair<int, int>>& bounds) {
31
  colors_map_ = colors_map;
32
  bounds_ = bounds;
33
  ncolors_ = bounds.size();
34
  current_colorcounts_ = std::vector<int>(ncolors_, 0);
35
}
36

37
void FairnessConstraint::Reset() {
38
  std::fill(current_colorcounts_.begin(), current_colorcounts_.end(), 0);
39
  current_set_.clear();
40
}
41

42
bool FairnessConstraint::CanAdd(int element) const {
43
  assert(!current_set_.count(element));
44
  int elt_color = colors_map_.at(element);
45
  return current_colorcounts_[elt_color] + 1 <= bounds_[elt_color].second;
46
}
47

48
void FairnessConstraint::Add(int element) {
49
  assert(!current_set_.count(element));
50
  int elt_color = colors_map_.at(element);
51
  current_colorcounts_[elt_color]++;
52
  current_set_.insert(element);
53
}
54

55
bool FairnessConstraint::CanRemove(int element) const {
56
  assert(current_set_.count(element));
57
  int elt_color = colors_map_.at(element);
58
  return current_colorcounts_[elt_color] - 1 >= bounds_[elt_color].first;
59
}
60

61
void FairnessConstraint::Remove(int element) {
62
  assert(current_set_.count(element));
63
  int elt_color = colors_map_.at(element);
64
  current_colorcounts_[elt_color]--;
65
  current_set_.erase(element);
66
}
67

68
int FairnessConstraint::GetColor(int element) const {
69
  return colors_map_.at(element);
70
}
71

72
int FairnessConstraint::GetColorNum() const { return bounds_.size(); }
73

74
std::vector<std::pair<int, int>> FairnessConstraint::GetBounds() const {
75
  return bounds_;
76
}
77

78
bool FairnessConstraint::IsFeasible(const std::vector<int> elements) {
79
  std::vector<int> colorcounts = std::vector<int>(ncolors_, 0);
80
  int elt_color;
81
  for (int elt : elements) {
82
    elt_color = colors_map_.at(elt);
83
    colorcounts[elt_color]++;
84
    if (colorcounts[elt_color] > bounds_[elt_color].second) return false;
85
  }
86
  for (int color = 0; color < ncolors_; color++) {
87
    if (colorcounts[color] < bounds_[color].first) return false;
88
  }
89
  return true;
90
}
91

92
std::unique_ptr<Matroid> FairnessConstraint::LowerBoundsToMatroid() const {
93
  std::vector<int> ks;
94
  ks.reserve(bounds_.size());
95
  for (const std::pair<int, int>& bound : bounds_) {
96
    ks.push_back(bound.first);
97
  }
98
  return std::make_unique<PartitionMatroid>(colors_map_, ks);
99
}
100

101
std::unique_ptr<Matroid> FairnessConstraint::UpperBoundsToMatroid() const {
102
  std::vector<int> ks;
103
  ks.reserve(bounds_.size());
104
  for (const std::pair<int, int>& bound : bounds_) {
105
    ks.push_back(bound.second);
106
  }
107
  return std::make_unique<PartitionMatroid>(colors_map_, ks);
108
}
109

110
std::unique_ptr<FairnessConstraint> FairnessConstraint::Clone() const {
111
  return std::make_unique<FairnessConstraint>(*this);
112
}
113