Adding multiple values to a key in a TreeMap in Java 8

Question

I have an <Integer, <Set<Element>>> TreeMap that I'm trying to add keys and values to, however I don't know how to add multiple values to keys, I've looked at other posts and searched online but couldn't find anything for my code. I have a list of elements (periodic elements) of which I'm iterating over, the elements contain details of element name, group number, weight etc. and I'm trying to use the group number as the key, and then the use all the other values (weight, name etc.), including the key as they values.

So something like this:

1(key):
1,H,Hydrogen,1(group num),1,1.008
3,Li,Lithium,1,2,6.94
11,Na,Sodium,1,3,22.98976928
2(key):
4,Be,Beryllium,2(group num),2,9.0121831
12,Mg,Magnesium,2,3,24.305
20,Ca,Calcium,2,4,40.078

I'm using a for loop to iterate through the list, and then I check if the group number is a key in the map, if it is I add the element to a set. If the group number is not in the map, then I add the element to a new set, then add that group number and new set to the map.

My code:

public static Map<Integer, Set<Element>> elementsByGroup(List<Element> elements){
Map<Integer, Set<Element>> elementMap = new TreeMap<>();
Set<Element> elementSet = new TreeSet<>(new ElementNumComparator());

for (Element i: elements){
    if(elementMap.containsKey(i.group)){
        elementSet.add(i);

    }
    else{
       Set<Element> newSet = new TreeSet<>(new ElementNumComparator());
       newSet.add(i);
       elementMap.put(i.group, newSet);
    }
}

Which outputs this when printed:

1/[55, Cs, Caesium, 1, 6, 132.905452]
2/[56, Ba, Barium, 2, 6, 137.327]
3/[89, Ac, Actinium, 3, 7, 227.0]
4/[72, Hf, Hafnium, 4, 6, 178.49]
5/[41, Nb, Niobium, 5, 5, 92.90637]
6/[24, Cr, Chromium, 6, 4, 51.9961]
7/[25, Mn, Manganese, 7, 4, 54.938044]
8/[26, Fe, Iron, 8, 4, 55.845]
9/[27, Co, Cobalt, 9, 4, 58.933194]
10/[28, Ni, Nickel, 10, 4, 58.6934]
11/[29, Cu, Copper, 11, 4, 63.546]
12/[48, Cd, Cadmium, 12, 5, 112.414]
13/[13, Al, Aluminium, 13, 3, 26.9815385]
14/[6, C, Carbon, 14, 2, 12.011]
15/[51, Sb, Antimony, 15, 5, 121.76]
16/[8, O, Oxygen, 16, 2, 15.999]
17/[85, At, Astatine, 17, 6, 210.0]
18/[18, Ar, Argon, 18, 3, 39.948]

However it only prints one of the elements per group, when there should be multiple. I tried adding the elements and group in the if statement but it just added every element to each key. I'd appreciate any help on this.

Answer 1

The problem is that you have a single elementSet set, that's not in your map, that you're adding elements to. The first time you see a given key, you're adding the element to the right place; but after that, any subsequent elements with the same key end up in elementSet , never to be seen again.

To fix this, you need to eliminate elementSet , and instead use elementMap.get(i.group) to find the set to add the element to.

Answer 2

The algorithm that you described is correct. Your implementation has a small issue: the check is done correctly, but you are adding to a wrong set -

if(elementMap.containsKey(i.group)){
    elementSet.add(i);
}

should be

if(elementMap.containsKey(i.group)){
    elementMap.get(i.group).add(i);
}

You could unify code paths a little more by changing the code as follows:

Set<Element> current = elementMap.get(i.group);
if (current == null) {
    current = new TreeSet<>(new ElementNumComparator());
    elementMap.put(i.group, current);
}
current.add(i);

This combines the check if a key is present or not with retrieval of the corresponding Set<Element> , and has a single call of add(i) for both code branches.