com.twitter.common.collections.Multimaps.java Source code

Introduction

Here is the source code for com.twitter.common.collections.Multimaps.java

Source

// =================================================================================================
// Copyright 2011 Twitter, Inc.
// -------------------------------------------------------------------------------------------------
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this work except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file, or 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 com.twitter.common.collections;

import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multiset;
import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;

/**
 * Utility class for functions related to Multimaps in the java collections library.
 *
 * @author William Farner
 */
public final class Multimaps {

    private Multimaps() {
        // Utility.
    }

    /**
     * Prunes a multimap based on a predicate, returning the pruned values.  The input map will be
     * modified.
     *
     * @param map The multimap to prune.
     * @param filterRule The pruning rule.  When the predicate returns {@code false} for an entry, it
     *    will be pruned, otherwise it will be retained.
     * @param <K> The key type in the multimap.
     * @param <V> The value type in the multimap.
     * @return A new multimap, containing the pruned keys/values.
     */
    public static <K, V> Multimap<K, V> prune(Multimap<K, V> map, Predicate<? super Collection<V>> filterRule) {
        Preconditions.checkNotNull(map);
        Preconditions.checkNotNull(filterRule);
        Multimap<K, V> pruned = ArrayListMultimap.create();
        Iterator<Map.Entry<K, Collection<V>>> asMapItr = map.asMap().entrySet().iterator();
        while (asMapItr.hasNext()) {
            Map.Entry<K, Collection<V>> asMapEntry = asMapItr.next();
            if (!filterRule.apply(asMapEntry.getValue())) {
                pruned.putAll(asMapEntry.getKey(), asMapEntry.getValue());
                asMapItr.remove();
            }
        }

        return pruned;
    }

    private static final class AtLeastSize implements Predicate<Collection<?>> {
        private final int minSize;

        AtLeastSize(int minSize) {
            Preconditions.checkArgument(minSize >= 0);
            this.minSize = minSize;
        }

        @Override
        public boolean apply(Collection<?> c) {
            return c.size() >= minSize;
        }
    }

    /**
     * Convenience method to prune key/values pairs where the size of the value collection is below a
     * threshold.
     *
     * @param map The multimap to prune.
     * @param minSize The minimum size for retained value collections.
     * @param <K> The key type in the multimap.
     * @param <V> The value type in the multimap.
     * @return A new multimap, containing the pruned keys/values.
     * @throws IllegalArgumentException if minSize < 0
     */
    public static <K, V> Multimap<K, V> prune(Multimap<K, V> map, int minSize) {
        return prune(map, new AtLeastSize(minSize));
    }

    /**
     * Returns the set of keys associated with groups of a size greater than or equal to a given size.
     *
     * @param map The multimap to search.
     * @param minSize The minimum size to return associated keys for.
     * @param <K> The key type for the multimap.
     * @return The keys associated with groups of size greater than or equal to {@code minSize}.
     * @throws IllegalArgumentException if minSize < 0
     */
    public static <K> Set<K> getLargeGroups(Multimap<K, ?> map, int minSize) {
        return Sets.newHashSet(Maps.filterValues(map.asMap(), new AtLeastSize(minSize)).keySet());
    }

    /**
     * Returns the set of keys associated with the largest values in the multimap.
     *
     * @param map The multimap to search.
     * @param topValues Number of groupings to find the keys for.
     * @return The keys associated with the largest groups, of maximum size {@code topValues}.
     */
    public static <K> Set<K> getLargestGroups(Multimap<K, ?> map, int topValues) {
        Ordering<Multiset.Entry<K>> groupOrdering = new Ordering<Multiset.Entry<K>>() {
            @Override
            public int compare(Multiset.Entry<K> entry1, Multiset.Entry<K> entry2) {
                return entry1.getCount() - entry2.getCount();
                // overflow-safe, since sizes are nonnegative
            }
        };
        Set<K> topKeys = Sets.newHashSetWithExpectedSize(topValues);
        for (Multiset.Entry<K> entry : groupOrdering.greatestOf(map.keys().entrySet(), topValues)) {
            topKeys.add(entry.getElement());
        }
        return topKeys;
    }
}