lang_2023/include/utils.hpp

#pragma once

#include <vector>
#include <unordered_map>

namespace utils {

using std::size_t;

using IdType = size_t;

enum class ReferenceType { Reference, UniqueReference };

template<typename T>
class Storage {
public:
  Storage() = default;

  IdType GetId(const T& value) {
    IdType id = 0;
    auto value_position = value_to_id_.find(value);

    if (value_position == value_to_id_.end()) {
      id = id_to_value_.size();
      value_to_id_[value] = id;
      id_to_value_.push_back(value);
    } else {
      id = value_position->second;
    }

    return id;
  }

  const T& GetValue(IdType id) {
    return id_to_value_[id];
  }

private:
  std::vector<T> id_to_value_;
  std::unordered_map<T, IdType> value_to_id_;
};

class GroupsManager { // TODO: recall right algorithm name
public:
  GroupsManager() = default;

  explicit GroupsManager(size_t n) {
    edges_.resize(n);
    ranks_.resize(n);

    for (size_t i = 0; i < n; ++i) {
      edges_[i] = i;
      ranks_[i] = 1;
    }
  }

  void Unite(size_t u, size_t v) { // TODO: recall choice algorithm
    u = GetGroupRoot(u);
    v = GetGroupRoot(v);
    if (ranks_[v] >= ranks_[u]) {
      edges_[u] = v;
      ranks_[v] = std::max(ranks_[u] + 1, ranks_[v]);
    } else {
      edges_[v] = u;
      // always ranks_[u] > ranks_[v]
    }
  }

  bool IsInOneGroup(size_t u, size_t v) {
    return GetGroupRoot(u) == GetGroupRoot(v);
  }

  size_t AddElement() {
    size_t id = edges_.size();
    edges_.push_back(id);
    ranks_.push_back(1);
    return id;
  }

  size_t GetGroupRoot(size_t v) {
    if (edges_[v] == v) {
      return v;
    }
    return edges_[v] = GetGroupRoot(edges_[v]);
  }
private:
  std::vector<size_t> edges_;
  std::vector<size_t> ranks_;
};

} // namespace utils