dynamic_connectivity.cpp

#include "dynamic_connectivity.hpp"
#include <algorithm>
#include <iomanip>
#include <iostream>
#include <omp.h>

DynamicConnectivity::DynamicConnectivity() {
}

DynamicConnectivity::DynamicConnectivity(long num_nodes)
    : num_nodes(num_nodes) {
  union_find = new std::atomic<Node>[num_nodes];
  for (int i = 0; i < num_nodes; ++i) {
    union_find[i] = i;
  }

  component_size = new std::atomic<uint>[num_nodes];
  for (int i = 0; i < num_nodes; ++i) {
    component_size[i] = 1;
  }

  // both directions
  forest_edges.reserve(2 * num_nodes);
  parent = new Node[num_nodes];
  for (int i = 0; i < num_nodes; ++i) {
    parent[i] = -2;
  }

  // locks.resize(num_nodes);
  locks = new std::mutex[num_nodes];
}

DynamicConnectivity::~DynamicConnectivity() {
  delete[] union_find;
  delete[] parent;
  delete[] component_size;
  delete[] locks;
}

void DynamicConnectivity::printUF() {
  for (int i = 0; i < 9; ++i) {
    std::cout << i << ": " << union_find[i] << std::endl;
  }
  std::cout << std::endl;
}

#define NOT_EMPTY head != queue.size()

#define PUSH(N, DIST) queue.push_back({N, DIST});

#define POP queue[head++]

struct QueueNode {
  long id;
  size_t distance;
};

void DynamicConnectivity::BFS(const Node s,
                              const SimpleAdjacencyArray<long> &graph) {
  // std::vector<char> visited(num_nodes, false);
  std::vector<QueueNode> queue;
  size_t head = 0;
  queue.reserve(num_nodes);
  PUSH(s, 0)
  while (NOT_EMPTY) {
    QueueNode queueNode = POP;
    SimpleAdjacencyArray<long>::NodeHandle currentHandle =
        graph.node(queueNode.id);
    for (auto nIt = graph.beginEdges(currentHandle);
         nIt != graph.endEdges(currentHandle); ++nIt) {
      long nID = graph.nodeId(graph.edgeHead(nIt));
      if (parent[nID] == -1) {
        parent[nID] = currentHandle;
        PUSH(nID, queueNode.distance + 1);
      }
    }
  }
}

void DynamicConnectivity::addEdges(const EdgeList &edges) {
#pragma omp parallel
  {
    EdgeList local_edges;
    local_edges.reserve(edges.size() / omp_get_num_threads() + 1);
#pragma omp for nowait
    for (uint i = 0; i < edges.size(); ++i) {
      // std::cout << std::setw(2) << edges[i].from << " | " << std::setw(2) <<
      // edges[i].to
      //           << std::endl;

      // find bigger component
      Node root_from = find_representative(edges[i].from);
      Node root_to = find_representative(edges[i].to);

      // same component
      if (root_from == root_to) {
        continue;
      }

      // avoids deadlock, alternative is to use order of locks index
      std::lock(locks[root_from], locks[root_to]);
      // has changed in the mean time
      while (union_find[root_to] != root_to ||
             union_find[root_from] != root_from) {
        // std::cout << i << " error: " << root_from << " " << root_to <<
        // std::endl; release all
        locks[root_from].unlock();
        locks[root_to].unlock();
        root_from = find_representative(root_from);
        root_to = find_representative(root_to);
        if (root_from == root_to) {
          continue;
        }
        std::lock(locks[root_from], locks[root_to]);
      }

      // union by rank
      if (component_size[root_from] > component_size[root_to]) {
        union_find[root_to] = root_from;
        component_size[root_from] += component_size[root_to];
      } else {
        union_find[root_from] = root_to;
        component_size[root_to] += component_size[root_from];
      }
      // union_find[root_from] = root_to;

      locks[root_from].unlock();
      locks[root_to].unlock();

      local_edges.push_back(edges[i]);
    }
    // std::cout << "critical " << omp_get_thread_num() << std::endl;
#pragma omp critical
    {
      for (auto [from, to, length] : local_edges) {
        forest_edges.push_back({from, to, length});
        forest_edges.push_back({to, from, length});
      }
    }
  }

  // used as visited in BFS
#pragma omp parallel for
  for (int i = 0; i < num_nodes; ++i) {
    parent[i] = -2;
  }
  SimpleAdjacencyArray<long> forest(num_nodes, forest_edges);

  if (roots.empty()) {
    for (long i = 0; i < num_nodes; ++i) {
      if (union_find[i] == i) {
        roots.push_back(i);
      }
    }
  } else {
    roots.erase(
        std::remove_if(roots.begin(), roots.end(),
                       [this](Node root) { return union_find[root] != root; }),
        roots.end());
  }
#pragma omp parallel for
  for (int i = 0; i < roots.size(); ++i) {
    parent[roots[i]] = -1;
    BFS(roots[i], forest);
  }
}

bool DynamicConnectivity::connected(Node a, Node b) const {
  return find_representative(a) == find_representative(b);
}

DynamicConnectivity::Node DynamicConnectivity::parentOf(Node n) const {
  return parent[n];
}

size_t DynamicConnectivity::sizeOf(Node n) const {
  return component_size[n];
}

DynamicConnectivity::Node
DynamicConnectivity::find_representative(Node a) const {
  Node current = a;
  Node parent = union_find[current];

  while (current != parent) {
    current = parent;
    parent = union_find[current];
  }

  // path compression
  Node representative = parent;
  current = a;
  parent = union_find[current];
  while (current != parent) {
    // may fail, I don't care
    union_find[current].compare_exchange_weak(current, representative);
    // union_find[current] = representative;
    current = parent;
    parent = union_find[current];
  }

  return parent;
}