org.opentripplanner.common.DisjointSet Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of otp Show documentation

The OpenTripPlanner multimodal journey planning system

There is a newer version: 2.6.0

package org.opentripplanner.common;

import gnu.trove.list.TIntList;
import gnu.trove.list.array.TIntArrayList;
import gnu.trove.map.TLongObjectMap;
import gnu.trove.map.TObjectIntMap;
import gnu.trove.map.hash.TLongObjectHashMap;
import gnu.trove.map.hash.TObjectIntHashMap;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import org.opentripplanner.util.MapUtils;

/** Basic union-find data structure with path compression */
public class DisjointSet {

  TIntList sets = new TIntArrayList();

  TObjectIntMap setMapping = new TObjectIntHashMap<>();

  public DisjointSet() {}

  public int union(T element1, T element2) {
    int p1 = find(element1);
    int p2 = find(element2);

    if (p1 == p2) {
      return p1;
    }

    int p1size = -sets.get(p1);
    int p2size = -sets.get(p2);

    int totalSize = p1size + p2size;

    if (p1size > p2size) {
      sets.set(p2, p1);
      sets.set(p1, -totalSize);
      return p2;
    } else {
      sets.set(p1, p2);
      sets.set(p2, -totalSize);
      return p1;
    }
  }

  public int find(T element) {
    if (setMapping.containsKey(element)) {
      int i = setMapping.get(element);
      return compact(i);
    } else {
      setMapping.put(element, sets.size());
      sets.add(-1);
      return sets.size() - 1;
    }
  }

  public boolean exists(T element) {
    return setMapping.containsKey(element);
  }

  public List> sets() {
    TLongObjectMap> out = new TLongObjectHashMap<>();
    setMapping.forEachEntry((k, v) -> {
      MapUtils.addToMapSet(out, compact(v), k);
      return true;
    });
    return new ArrayList<>(out.valueCollection());
  }

  public int size(int component) {
    return -sets.get(component);
  }

  private int compact(int i) {
    int key = sets.get(i);
    if (key < 0) {
      return i;
    }
    int j = compact(key);
    sets.set(i, j);
    return j;
  }
}