xref: /aosp_15_r20/external/cronet/net/first_party_sets/addition_overlaps_union_find.cc (revision 6777b5387eb2ff775bb5750e3f5d96f37fb7352b)
1 // Copyright 2022 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #include "net/first_party_sets/addition_overlaps_union_find.h"
6 
7 #include <numeric>
8 
9 #include "base/check_op.h"
10 #include "base/containers/flat_map.h"
11 #include "base/containers/flat_set.h"
12 
13 namespace net {
14 
AdditionOverlapsUnionFind(int num_sets)15 AdditionOverlapsUnionFind::AdditionOverlapsUnionFind(int num_sets) {
16   CHECK_GE(num_sets, 0);
17   representatives_.resize(num_sets);
18   std::iota(representatives_.begin(), representatives_.end(), 0ul);
19 }
20 
21 AdditionOverlapsUnionFind::~AdditionOverlapsUnionFind() = default;
22 
Union(size_t set_x,size_t set_y)23 void AdditionOverlapsUnionFind::Union(size_t set_x, size_t set_y) {
24   CHECK_GE(set_x, 0ul);
25   CHECK_LT(set_x, representatives_.size());
26   CHECK_GE(set_y, 0ul);
27   CHECK_LT(set_y, representatives_.size());
28 
29   size_t root_x = Find(set_x);
30   size_t root_y = Find(set_y);
31 
32   if (root_x == root_y)
33     return;
34   auto [parent, child] = std::minmax(root_x, root_y);
35   representatives_[child] = parent;
36 }
37 
SetsMapping()38 AdditionOverlapsUnionFind::SetsMap AdditionOverlapsUnionFind::SetsMapping() {
39   SetsMap sets;
40 
41   // An insert into the flat_map and flat_set has O(n) complexity and
42   // populating sets this way will be O(n^2).
43   // This can be improved by creating an intermediate vector of pairs, each
44   // representing an entry in sets, and then constructing the map all at once.
45   // The intermediate vector stores pairs, using O(1) Insert. Another vector
46   // the size of |num_sets| will have to be used for O(1) Lookup into the
47   // first vector. This means making the intermediate vector will be O(n).
48   // After the intermediate vector is populated, and we can use
49   // base::MakeFlatMap to construct the mapping all at once.
50   // This improvement makes this method less straightforward however.
51   for (size_t i = 0; i < representatives_.size(); i++) {
52     size_t cur_rep = Find(i);
53     auto it = sets.emplace(cur_rep, base::flat_set<size_t>()).first;
54     if (i != cur_rep) {
55       it->second.insert(i);
56     }
57   }
58   return sets;
59 }
60 
Find(size_t set)61 size_t AdditionOverlapsUnionFind::Find(size_t set) {
62   CHECK_GE(set, 0ul);
63   CHECK_LT(set, representatives_.size());
64   if (representatives_[set] != set)
65     representatives_[set] = Find(representatives_[set]);
66   return representatives_[set];
67 }
68 
69 }  // namespace net
70