Add two collections together, merging if present, adding if not

Question

When adding two TVShows together, assuming they are the same show, if TVShow one has a Season 1 and TVShow two has a Season 1, it will not merge the episodes from Season 1 in TVShow two to Season 1 in TVShow one as it thinks they are equal. However this issue is that if I also added the episodes in the equals function of Season, then TVShow one would contain two entries, both being Season 1 rather than merging the episodes from two to one.

Is TreeSet the correct collection for this?
And how should I go about doing this?

public class TVShow extends Show implements Iterable<Season> {
    private final TreeSet<Season> seasons;
}

public class Season implements Iterable<Episode>, Comparable<Season> {
    private final TreeSet<Episode> episodes;
    private String name = null;
    private int number;

    @Override
    public boolean equals(Object o) {
        if (o == this)
            return true;
        else if (!(o instanceof Season))
            return false;
        else {
            Season other = (Season) o;
            return this.number == other.number;
        }
    }
}

public class Episode implements Comparable<Episode> {
    private String name;
    private int number;
}

For example:

TVShow one = new TVShow();
Season s1 = new Season(1);
s1.add(new Episode(1));
one.add(s1);

TVShow two = new TVShow();
Season sOne = new Season(1);
sOne.add(new Episode(1));
sOne.add(new Episode(2));
sOne.add(new Episode(3));
two.add(sOne);

one.add(two);

Result: one = {seasons{1, episodes{1, 2, 3}}}

Answer 1

Merging 2 seasons is simple enough, since the episodes field is a Set, you just call addAll() and any duplicates are ignored.

Merging 2 shows is more complex, because when a duplicate season is found, they need to be merged, so you can't just call addAll(), since that will add new seasons, but won't merge existing seasons.

The problem is that TreeSet doesn't have a method for getting the existing season by number. One way to fix that is it change the Set to a Map, keyed by number. That would be the recommended way, but there is a trick using subSet() and first() for getting the object in the Set that "equals" the object being looked up.

So, to complete the code, eliminating the name fields and the Iterable interfaces, as they are immaterial to the question, we get:

final class TVShow {
    private final TreeSet<Season> seasons = new TreeSet<>();

    public TVShow(Season... seasons) {
        this.seasons.addAll(Arrays.asList(seasons));
    }
    public void merge(TVShow that) {
        for (Season season : that.seasons) {
            // Add if season is new, otherwise merge it
            if (! this.seasons.add(season)) {
                this.seasons.subSet(season, true, season, true)
                            .first().merge(season);
            }
        }
    }
    @Override
    public String toString() {
        return this.seasons.toString();
    }
}

final class Season implements Comparable<Season> {
    private final int number;
    private final TreeSet<Episode> episodes = new TreeSet<>();

    public Season(int number, Episode... episodes) {
        this.number = number;
        this.episodes.addAll(Arrays.asList(episodes));
    }
    @Override
    public int compareTo(Season that) {
        return Integer.compare(this.number, that.number);
    }
    public void merge(Season that) {
        // Just add missing episodes
        this.episodes.addAll(that.episodes);
    }
    @Override
    public String toString() {
        return this.number + ": " + this.episodes;
    }
}

final class Episode implements Comparable<Episode> {
    private final int number;

    public Episode(int number) {
        this.number = number;
    }
    @Override
    public int compareTo(Episode that) {
        return Integer.compare(this.number, that.number);
    }
    @Override
    public String toString() {
        return String.valueOf(this.number);
    }
}

Test

TVShow show1 = new TVShow(
        new Season(1, new Episode(1), new Episode(2), new Episode(3)),
        new Season(2, new Episode(1), new Episode(2), new Episode(3)));
TVShow show2 = new TVShow(
        new Season(1, new Episode(2), new Episode(4), new Episode(6)),
        new Season(3, new Episode(1), new Episode(2), new Episode(4)));
System.out.println("show1: " + show1);
System.out.println("show2: " + show2);
System.out.println();

show1.merge(show2);
System.out.println("show1: " + show1);
System.out.println("show2: " + show2);

Output

show1: [1: [1, 2, 3], 2: [1, 2, 3]]
show2: [1: [2, 4, 6], 3: [1, 2, 4]]

show1: [1: [1, 2, 3, 4, 6], 2: [1, 2, 3], 3: [1, 2, 4]]
show2: [1: [2, 4, 6], 3: [1, 2, 4]]

Output shows that show2 is unmodified, and that in show1, season 1 was merged and season 3 was added.

---

UPDATE: Alternate version using Map instead of Set. Recommended!

final class TVShow {
    private final TreeMap<Integer, Season> seasons = new TreeMap<>();

    public TVShow(Season... seasons) {
        for (Season season : seasons)
            this.seasons.put(season.getNumber(), season);
    }
    public TVShow merge(TVShow that) {
        for (Season season : that.seasons.values())
            this.seasons.merge(season.getNumber(), season, Season::merge);
        return this;
    }
    @Override
    public String toString() {
        return this.seasons.values().toString();
    }
}

final class Season {
    private final int number;
    private final TreeMap<Integer, Episode> episodes = new TreeMap<>();

    public Season(int number, Episode... episodes) {
        this.number = number;
        for (Episode episode : episodes)
            this.episodes.put(episode.getNumber(), episode);
    }
    public int getNumber() {
        return this.number;
    }
    public Season merge(Season that) {
        for (Episode episode : that.episodes.values())
            this.episodes.merge(episode.getNumber(), episode, Episode::merge);
        return this;
    }
    @Override
    public String toString() {
        return this.number + ": " + this.episodes.values();
    }
}

final class Episode {
    private final int number;

    public Episode(int number) {
        this.number = number;
    }
    public int getNumber() {
        return this.number;
    }
    public Episode merge(Episode that) {
        // Nothing to merge
        return this;
    }
    @Override
    public String toString() {
        return String.valueOf(this.number);
    }
}