Search code examples
javasortingtreemap

Sort TreeMap based on a Enum


I've a TreeMap:

TreeMap<String, List<MyObject>> myTreeMap = new TreeMap<String, List<MyObject>>(new MyComparator());

I want to sort this TreeMap based on order of an enum.

MyActionEnum:

public enum MyActionEnum implements Serializable {

    ADD("Add"),
    UPDATE("Update"),
    DELETE("Delete");

    private String action;
    private static final long serialVersionUID = 1L;

    MyActionEnum(String action) {
        this.action = action;
    }

    public String get() {
        return this.action;
    }

}

Such that in my TreeMap, order should be Add, Update, Delete.

MyComparator:

public class MyComparator implements Comparator<String> {

    @Override
    public int compare(String o1, String o2) {
        // how to compare here based on enum.
    }
}

Solution

  • First, you need a way to lookup the enum value from the key. Best way to do that is a method on the enum itself, so it's reusable:

    public static MyActionEnum fromAction(String action) {
        for (MyActionEnum e : values())
            if (e.action.equals(action))
                return e;
        if (action == null)
            throw new NullPointerException("Action is null");
        throw new IllegalArgumentException("Unknown action: " + action);
    }
    

    If there were a lot of enum values, you'd want to replace the sequential search done here with a Map lookup.

    You can then use that in the compare() method:

    @Override
    public int compare(String o1, String o2) {
        return Integer.compare(MyActionEnum.fromAction(o1).ordinal(),
                               MyActionEnum.fromAction(o2).ordinal());
    }