第三大的数

给你一个非空数组，返回此数组中 第三大的数 。如果不存在，则返回数组中最大的数。

input：[3, 2, 1]
output：1

class Solution {
    public int thirdMax(int[] nums) {
        TreeSet<Integer> set = new TreeSet<Integer>();
        for(int num : nums){
            set.add(num);
            if(set.size() > 3){
                set.remove(set.first());
            }
        }
        return set.size() == 3 ? set.first() : set.last();
    }
}