SetOperations.java

package net.morimekta.collect.util;

import net.morimekta.collect.SetBuilder;
import net.morimekta.collect.UnmodifiableSet;

import java.util.Comparator;
import java.util.Objects;
import java.util.Set;
import java.util.SortedSet;

import static net.morimekta.collect.UnmodifiableSet.asSet;
import static net.morimekta.collect.UnmodifiableSet.setOf;
import static net.morimekta.collect.UnmodifiableSortedSet.asSortedSet;
import static net.morimekta.collect.UnmodifiableSortedSet.newBuilderOrderedBy;

/**
 * Utility class containing set-math operations.
 * <p>
 * See <a href="https://www.probabilitycourse.com/chapter1/1_2_2_set_operations.php">probabilitycouse</a>
 * or <a href="https://www.britannica.com/science/set-theory/Operations-on-sets">britannica</a>
 * for definitions and theory.
 */
public final class SetOperations {
    /**
     * Get the union of all the set's provided. Essentially a set
     * of all elements in each set in the input.
     *
     * @param of  Sets to make union of.
     * @param <E> The item type.
     * @return The union set.
     */
    @SafeVarargs
    public static <E> Set<E> union(Set<E>... of) {
        Objects.requireNonNull(of);
        if (of.length == 0) {
            return setOf();
        }
        SetBuilder<E> builder;
        if (of[0] instanceof SortedSet) {
            SortedSet<E> first = (SortedSet<E>) of[0];
            builder = newBuilderOrderedBy(first.comparator());
        } else {
            builder = UnmodifiableSet.newBuilder();
        }
        for (Set<E> set : of) {
            builder.addAll(set);
        }
        return builder.build();
    }

    /**
     * Get the intersection between the sets. These are the items contains
     * in all the sets provided.
     *
     * @param set  The main set to make intersection with.
     * @param with Sets to intersect with.
     * @param <E>  The item type.
     * @return The intersected set.
     */
    @SafeVarargs
    public static <E> Set<E> intersect(Set<E> set, Set<E>... with) {
        if (with.length == 0) {
            if (set instanceof SortedSet) {
                return asSortedSet(set);
            }
            return asSet(set);
        }
        SetBuilder<E> builder;
        if (set instanceof SortedSet) {
            SortedSet<E> first = (SortedSet<E>) set;
            builder = newBuilderOrderedBy(set.size(), first.comparator());
        } else {
            builder = UnmodifiableSet.newBuilder(set.size());
        }
        setLoop:
        for (E e : set) {
            for (Set<E> w : with) {
                if (!w.contains(e)) {
                    continue setLoop;
                }
            }
            builder.add(e);
        }
        return builder.build();
    }

    /**
     * Take the minuend, and subtract all elements contained in all the subtrahend sets.
     *
     * @param minuend     The source set to be removed elements from.
     * @param subtrahends The sets with elements to remove.
     * @param <E>         The element type.
     * @return The difference set from the subtraction.
     */
    @SafeVarargs
    @SuppressWarnings("unchecked,rawtypes")
    public static <E> Set<E> subtract(Set<E> minuend, Set<E>... subtrahends) {
        if (subtrahends.length == 0) {
            if (minuend instanceof SortedSet) {
                return asSortedSet(minuend);
            }
            return asSet(minuend);
        }
        SetBuilder<E> builder;
        if (minuend instanceof SortedSet) {
            SortedSet<E> first = (SortedSet<E>) minuend;
            builder = newBuilderOrderedBy((Comparator) first.comparator());
        } else {
            builder = UnmodifiableSet.newBuilder();
        }

        setLoop:
        for (E e : minuend) {
            for (Set<E> w : subtrahends) {
                if (w.contains(e)) {
                    continue setLoop;
                }
            }
            builder.add(e);
        }
        return builder.build();
    }

    /**
     * Create a difference set of the items contained only in either
     * of the two sets, but not both. This is the same as
     * <code>union(subtract(a, b), subtract(b, a))</code>, but much more
     * efficient.
     *
     * @param first  The first set.
     * @param second The second set.
     * @param <E>    The item type.
     * @return The difference set.
     */
    public static <E> Set<E> difference(Set<E> first, Set<E> second) {
        SetBuilder<E> builder;
        if (first instanceof SortedSet) {
            if (first.isEmpty()) {
                return asSortedSet(second);
            }
            if (second.isEmpty()) {
                return asSortedSet(first);
            }
            builder = newBuilderOrderedBy(
                    first.size() + second.size(),
                    ((SortedSet<E>) first).comparator());
        } else {
            if (first.isEmpty()) {
                return asSet(second);
            }
            if (second.isEmpty()) {
                return asSet(first);
            }
            builder = UnmodifiableSet.newBuilder(first.size() + second.size());
        }

        for (E e : first) {
            if (!second.contains(e)) {
                builder.add(e);
            }
        }
        for (E e : second) {
            if (!first.contains(e)) {
                builder.add(e);
            }
        }
        return builder.build();
    }

    /**
     * Calculate the cartesian product of the two sets. The cartesian product will
     * combine every element of A with every element of B, so the resulting set
     * will have a total of <code>size(A) * size(B)</code> elements each with a
     * unique combination of <code>(Ax, By)</code> pair.
     *
     * @param a   The first set
     * @param b   The second set
     * @param <A> The element type in first set
     * @param <B> The element type in second set
     * @return The cartesian product.
     */
    @SuppressWarnings("unchecked,rawtypes")
    public static <A, B> Set<Pair<A, B>> product(Set<A> a, Set<B> b) {
        SetBuilder<Pair<A, B>> builder;
        if (a instanceof SortedSet && b instanceof SortedSet) {
            SortedSet<A> sa = (SortedSet<A>) a;
            SortedSet<B> sb = (SortedSet<B>) b;

            Comparator<Pair<A, B>> comparator = sa.comparator() == null ?
                                                Comparator.comparing(p -> (Comparable) p.first) :
                                                Comparator.comparing(Pair::getFirst, sa.comparator());
            if (sb.comparator() == null) {
                comparator = comparator.thenComparing(p -> (Comparable) p.second);
            } else {
                comparator = comparator.thenComparing(p -> p.second, sb.comparator());
            }

            if (a.isEmpty() || b.isEmpty()) {
                return newBuilderOrderedBy(comparator).build();
            }

            builder = newBuilderOrderedBy(a.size() * b.size(), comparator);
        } else {
            if (a.isEmpty() || b.isEmpty()) {
                return UnmodifiableSet.setOf();
            }
            builder = UnmodifiableSet.newBuilder(a.size() * b.size());
        }

        for (A ia : a) {
            for (B ib : b) {
                builder.add(Pair.pairOf(ia, ib));
            }
        }
        return builder.build();
    }

    /**
     * Check if the first set is a subset of the second.
     *
     * @param set The set to check.
     * @param of  The set to check against.
     * @param <E> The item type.
     * @return True if the first set is a subset of the second.
     */
    public static <E> boolean subset(Set<E> set, Set<E> of) {
        if (set.size() > of.size()) {
            return false;
        }
        for (E e : set) {
            if (!of.contains(e)) {
                return false;
            }
        }
        return true;
    }

    /**
     * Check that the two sets are disjoint, meaning they have no items in common.
     * Same that the intersection between the sets are empty.
     *
     * @param first  The first set.
     * @param second The second set.
     * @param <E>    The item type.
     * @return True if the two sets are disjoint.
     */
    public static <E> boolean disjoint(Set<E> first, Set<E> second) {
        // We actually only need to check one way, as the equality
        // items always comes in pairs...
        for (E item : first) {
            if (second.contains(item)) {
                return false;
            }
        }
        return true;
    }

    private SetOperations() {}
}