One possible solution for Multiset test

HashMultiset.java

package com.example;

import org.jetbrains.annotations.NotNull;

import java.util.*;

/**
 * Implements multiset based on LinkedHashMap class.
 * @param <E> Type of stored elements.
 */
public class HashMultiset<E> extends AbstractSet<E> implements Multiset<E> {

    /** Underlying {@see LinkedHashMap} which allows to meet asymptotic complexity requirements. */
    private final Map<E, Integer> underlyingSet = new LinkedHashMap<>();

    /** Overall quantity of elements in multiset. Can be calculated, but is cached here to meet O(1) constraint. */
    private int size = 0;

    /** {@inheritDoc} */
    @Override
    public int count(Object element) {
        final Integer result = underlyingSet.get(element);
        return result == null ? 0 : result;
    }

    /** {@inheritDoc} */
    @Override
    public boolean add(E element) {
        final Integer previous = underlyingSet.putIfAbsent(element, 1);
        if (previous != null) {
            underlyingSet.put(element, previous + 1);
        }

        ++size;
        return true;
    }

    /** {@inheritDoc} */
    @Override
    public boolean contains(Object element) {
        return underlyingSet.containsKey(element);
    }

    /** {@inheritDoc} */
    @Override
    public boolean remove(Object element) {
        final Integer count = underlyingSet.get(element);
        if (count == null) {
            return false;
        }

        if (count == 1) {
            underlyingSet.remove(element);
        } else {
            underlyingSet.put((E) element, count - 1);
        }

        size--;
        return true;
    }

    /** {@inheritDoc} */
    @Override
    public Set<E> elementSet() {
        return underlyingSet.keySet();
    }

    /** {@inheritDoc} */
    @Override
    public Set<Entry<E>> entrySet() {
        return new EntrySet();
    }

    /** {@inheritDoc} */
    @Override
    public Iterator<E> iterator() {
        return new MultisetIterator();
    }

    /** {@inheritDoc} */
    @Override
    public int size() {
        return size;
    }

    /** Iterator for multiset, returns all contained elements. Different elements returned in order of addition. */
    private class MultisetIterator implements Iterator<E>
    {
        private Iterator<Map.Entry<E, Integer>> underlyingIterator = underlyingSet.entrySet().iterator();
        private Map.Entry<E, Integer> currentEntry = null;
        private int currentCount = 0;
        private boolean wasRemove = false;

        /** {@inheritDoc} */
        @Override
        public boolean hasNext() {
            return underlyingIterator.hasNext()
                || (currentEntry != null && currentCount < currentEntry.getValue() - 1);
        }

        /** {@inheritDoc} */
        @Override
        public E next() {
            if (currentEntry != null && currentCount < currentEntry.getValue() - 1) {
                ++currentCount;
            } else {
                currentEntry = underlyingIterator.next();
                currentCount = 0;
            }

            wasRemove = false;
            return currentEntry.getKey();
        }

        /** {@inheritDoc} */
        @Override
        public void remove() {
            if (wasRemove) {
                throw new IllegalStateException();
            }

            currentEntry.setValue(currentEntry.getValue() - 1);
            currentCount--;

            if (currentEntry.getValue() == 0) {
                currentCount = 0;
                currentEntry = null;

                underlyingIterator.remove();
            }

            --size;
            wasRemove = true;
        }
    }

    /** Custom set for entrySet() method, needed to meet O(1) constraint. */
    private class EntrySet extends AbstractSet<Entry<E>> {

        /** {@inheritDoc} */
        @Override
        public Iterator iterator() {
            return new EntrySetIterator();
        }

        /** {@inheritDoc} */
        @Override
        public int size() {
            return underlyingSet.size();
        }

        private class EntrySetIterator implements Iterator<Entry<E>> {
            private final Iterator<Map.Entry<E, Integer>> underlyingIterator = underlyingSet.entrySet().iterator();
            private HashEntry currentEntry = null;

            /** {@inheritDoc} */
            @Override
            public boolean hasNext() {
                return underlyingIterator.hasNext();
            }

            /** {@inheritDoc} */
            @Override
            public Entry<E> next() {
                Map.Entry<E, Integer> entry = underlyingIterator.next();
                currentEntry = new HashEntry(entry.getKey(), entry.getValue());
                return currentEntry;
            }

            /** {@inheritDoc} */
            @Override
            public void remove() {
                underlyingIterator.remove();
                size -= currentEntry.getCount();
            }
        }

        /** Entry of a data structure. */
        private class HashEntry implements Multiset.Entry<E> {
            private final E element;
            private final Integer count;

            /**
             * Creates a new object of {@see HashEntry} class.
             * @param element stored multiset element.
             * @param count quantity of that element in multiset.
             */
            public HashEntry(E element, Integer count) {
                this.element = element;
                this.count = count;
            }

            /** {@inheritDoc} */
            @Override
            public @NotNull E getElement() {
                return element;
            }

            /** {@inheritDoc} */
            @Override
            public int getCount() {
                return count;
            }
        }
    }
}