EquivalenceTreeSet.java

/**
 * Copyright (C) 2008 Logan Johnson
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package loganj.structures;

import com.google.common.collect.*;
import com.google.common.base.Join;

import java.util.*;

/**
 * Set which differs from TreeSet only in that it allows multiple instances which are equivalent per their
 * natural ordering or the supplied Comparator.
 *
 * EquivalenceTreeSet contains no pair of elements e1 and e2 such that e1.equals(e2),
 * but may contain two elements e3 and e4 such that e3.compareTo(e4) == 0.
 * 
 * Equivalence groups are sorted by natural order or supplied Comparator.  Within an equivalence group, insertion order holds.
 *
 * Backed by a TreeMap (to store equivalence classes) and LinkedLists (to store equivalent elements within a class).
 *
 * @author Logan Johnson <logan.johnson@gmail.com>
 */
final public class EquivalenceTreeSet<T> extends AbstractSet<T> implements SortedSet<T> {

    final private SortedMap<T, List<T>> map;

    EquivalenceTreeSet() {
        this(new TreeMap<T,List<T>>());
    }

    EquivalenceTreeSet(Comparator<T> comparator) {
        this(new TreeMap<T,List<T>>(comparator));
    }

    private EquivalenceTreeSet(SortedMap<T, List<T>> backingMap) {
        this.map = backingMap;
    }

    @Override
    public int size() {
        // lame; just keep a counter.
        int size = 0;
        for ( Collection<T> e : map.values() ) {
            size += e.size();
        }
        return size;
    }

    @Override
    public boolean contains(Object o) {
        return map.containsKey(o) && map.get(o).contains(o);
    }
    
    @Override
    public Iterator<T> iterator() {
        return new Iterator<T>() {
            final private Iterator<T> mapIterator = map.keySet().iterator();
            private Iterator<T> equivalenceIterator = null;

            public boolean hasNext() {
                return ((equivalenceIterator != null) && equivalenceIterator.hasNext()) || mapIterator.hasNext();
            }

            public T next() {
                while ( (equivalenceIterator == null) || !equivalenceIterator.hasNext() ) {
                    equivalenceIterator = map.get(mapIterator.next()).iterator();
                }
                return equivalenceIterator.next();
            }

            public void remove() {
                equivalenceIterator.remove();
            }
        };
    }

    @Override
    public boolean add(T t) {
        List<T> e = map.get(t);
        if ( e == null ) {
            e = Lists.newLinkedList();
            map.put(t, e);
        }
        if ( e.contains(t) ) {
            return false;
        } else {
            e.add(t);
            return true;
        }
    }

    @Override
    public boolean remove(Object o) {
        if ( ! map.containsKey(o) ) {
            return false;
        }
        boolean removed = map.get(o).remove(o);
        if ( removed && map.get(o).isEmpty() ) {
            map.remove(o);
        }
        return removed;
    }

    @Override
    public void clear() {
        map.clear();
    }


    @Override
    public Comparator<? super T> comparator() {
        return map.comparator();
    }

    @Override
    public SortedSet<T> subSet(T t, T t1) {
        return new EquivalenceTreeSet<T>(map.subMap(t, t1));
    }

    @Override
    public SortedSet<T> headSet(T t) {
        return new EquivalenceTreeSet<T>(map.headMap(t));
    }

    @Override
    public SortedSet<T> tailSet(T t) {
        return new EquivalenceTreeSet<T>(map.tailMap(t));
    }

    @Override
    public T first() {
        return map.get(map.firstKey()).get(0);
    }

    @Override
    public T last() {
        List<T> list = map.get(map.lastKey());
        return list.get(list.size()-1);
    }

}