Map Inversion

MapInverter.java

package java8.util;

import static java.util.stream.Collectors.groupingBy;
import static java.util.stream.Collectors.mapping;
import static java.util.stream.Collectors.toCollection;
import static java.util.stream.Collectors.toMap;

import java.util.AbstractMap.SimpleEntry;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Stream;

/**
 * After reading below discussions on stackoverflow, I decided to create my own map inversion utility, internally using Java streams.
 * @author Sebastian Millies
 * @see "http://stackoverflow.com/questions/36795890/functionally-inverting-a-map-with-a-nested-data-structure-in-java-8"
 * @see "http://stackoverflow.com/questions/20412354/reverse-hashmap-keys-and-values-in-java"
 * @see "http://stackoverflow.com/questions/30021229/reverse-map-structure-using-java-8-streams"
 */
public class MapInverter {

    /**
     * Inverts a map. Would also work with multimaps if they are like Guava's multimaps, i. e. the entrySet contains only elements with
     * single values.
     * 
     * @param map the map to inverted
     * @param keyGenerator a function that converts an original map value to a key in the inverted map
     * @param mapFactory a factory which returns a new empty Map into which the results will be inserted
     * @param collectionFactory a factory for the value type of the inverted map
     * @return the inverted map, where all keys that map to the same value are now mapped from that value
     */
    public static <K, V, V1, C extends Collection<K>> Map<V, C> invertMap(Map<K, V1> map, Function<V1, ? extends V> keyGenerator,
            Supplier<Map<V, C>> mapFactory, Supplier<C> collectionFactory) {
        return invertMap(map.entrySet(), keyGenerator, mapFactory, collectionFactory);
    }

    /**
     * Inverts a map given an entry collection view of that map. The entry set may be that of a Guava multimap, i. e. contain duplicate keys
     * that are mapped to single values. Duplicate key value pairs are also supported, but will be mapped to a single entry in the inverted
     * map.
     * 
     * @param entries a map entry collection to be inverted
     * @param keyGenerator a function that converts an original map value to a key in the inverted map
     * @param mapFactory a factory which returns a new empty Map into which the results will be inserted
     * @param collectionFactory a factory for the value type of the inverted map
     * @return the inverted map, where all keys that map to the same value are now mapped from that value
     */
    public static <K, V, V1, C extends Collection<K>> Map<V, C> invertMap(Collection<Entry<K, V1>> entries, Function<V1, ? extends V> keyGenerator,
            Supplier<Map<V, C>> mapFactory, Supplier<C> collectionFactory) {
        Map<V, C> inverted = entries.stream().collect(groupingBy( e -> keyGenerator.apply(e.getValue()), mapFactory, 
                                                            mapping(Entry::getKey, toCollection(collectionFactory))));
        return inverted;
    }

    /**
     * Convenience method for the simplest and most common case of {@link #invertMap(Map)}. The inverted map is implemented by
     * {@code HashMap}, will contain the original values as keys, and will contain {@code HashSet}s of the original keys as values.
     */
    public static <K, V> Map<V, Set<K>> invertMap(Map<K, V> map) {
        return invertMap(map, Function.identity(), HashMap::new, HashSet::new);
    }
    
    /**
     * Inverts a map. The map may also be a non-Guava multimap (i. e. the entrySet elements may have collections as values).
     * 
     * @param map the map to inverted
     * @param valueStreamer a function that converts an original map value to a stream. (In case of a multimap, may stream the original
     *            value collection.) Can also perform any other transformation on the original values to make them suitable as keys.
     * @param mapFactory a factory which returns a new empty Map into which the results will be inserted
     * @param collectionFactory a factory for the value type of the inverted map
     * @return the inverted map, where all keys that map to the same value are now mapped from that value
     */
    public static <K, V, V1, C extends Collection<K>> Map<V, C> invertMultiMap(Map<K, V1> map, Function<V1, Stream<V>> valueStreamer,
            Supplier<Map<V, C>> mapFactory, Supplier<C> collectionFactory) {
        Map<V, C> inverted = map.entrySet().stream().flatMap(e ->
                    valueStreamer.apply(e.getValue()).map(v -> new SimpleEntry<>(v, newCollection(e.getKey(), collectionFactory))))
                .collect(toMap(Entry::getKey, Entry::getValue, (a, b) -> {a.addAll(b);  return a;}, mapFactory));
        return inverted;
    }

    private static <T, E extends T, C extends Collection<T>> C newCollection(E elem, Supplier<C> s) {
        C collection = s.get();
        collection.add(elem);
        return collection;
    }

}

MapInverterTest.java

package java8.util;

import static com.google.common.collect.Sets.newHashSet;
import static java.util.Arrays.asList;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.function.Function;

import org.junit.Test;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Multimap;

public class MapInverterTest {

    @Test
    public void invertSimple() {
        Map<Long, Set<String>> inverted = MapInverter.invertMap(newSimpleMap());
        assertEquals(invertedSimpleMap(), inverted);
        assertTrue(inverted instanceof HashMap);
        inverted.values().forEach(s -> assertTrue(s instanceof HashSet));
    }

    @Test
    public void invertMulti() {
        Map<Long, Set<String>> inverted = MapInverter.invertMultiMap(newMultiMap(), Set::stream, TreeMap::new, TreeSet::new);
        assertEquals(invertedMultiMapSet(), inverted);
        assertTrue(inverted instanceof TreeMap);
        inverted.values().forEach(s -> assertTrue(s instanceof TreeSet));
    }
    
    @Test
    public void invertGuavaMulti() {
        Collection<Entry<String, Long>> entries = newGuavaMultiMap().entries();
        Map<Long, SortedSet<String>> inverted = MapInverter.invertMap(entries, Function.identity(), HashMap::new, TreeSet::new);
        assertEquals(invertedGuavaMultiMap(), inverted);
        assertTrue(inverted instanceof HashMap);
        inverted.values().forEach(s -> assertTrue(s instanceof TreeSet));
    }
    
    @Test
    public void invertMultiList() {
        Map<Long, List<String>> inverted = MapInverter.invertMultiMap(newMultiMap(), Set::stream, TreeMap::new, ArrayList::new);
        assertEquals(invertedMultiMapList(), inverted);
        assertTrue(inverted instanceof TreeMap);
        inverted.values().forEach(s -> assertTrue(s instanceof ArrayList));
    }

    private Map<String, Long> newSimpleMap() {
        Map<String, Long> map = new HashMap<>();
        map.put("a", 1L);
        map.put("b", 2L);
        map.put("c", 1L);
        return map;
    }
    
    private Multimap<String, Long> newGuavaMultiMap() {
        ListMultimap<String, Long> map = ArrayListMultimap.create();
        map.put("a", 1L);
        map.put("a", 1L);
        map.put("a", 2L);
        map.put("b", 2L);
        map.put("c", 1L);
        return map;
    }
    
    private Map<Long, Set<String>> invertedGuavaMultiMap() {
        Map<Long, Set<String>> map = new HashMap<>();
        map.put(1L, newHashSet("a", "c"));
        map.put(2L, newHashSet("a", "b"));
        return map;
    }

    private Map<Long, Set<String>> invertedSimpleMap() {
        Map<Long, Set<String>> map = new HashMap<>();
        map.put(1L, newHashSet("a", "c"));
        map.put(2L, newHashSet("b"));
        return map;
    }

    private Map<String, Set<Long>> newMultiMap() {
        Map<String, Set<Long>> map = new HashMap<>();
        map.put("a", newHashSet(1L, 2L));
        map.put("b", newHashSet(1L, 3L));
        return map;
    }

    private Map<Long, Set<String>> invertedMultiMapSet() {
        Map<Long, Set<String>> map = new HashMap<>();
        map.put(1L, newHashSet("a", "b"));
        map.put(2L, newHashSet("a"));
        map.put(3L, newHashSet("b"));
        return map;
    }
    
    private Map<Long, List<String>> invertedMultiMapList() {
        Map<Long, List<String>> map = new HashMap<>();
        map.put(1L, asList("a", "b"));
        map.put(2L, asList("a"));
        map.put(3L, asList("b"));
        return map;
    }

}