K-diff Pairs in an Array

https://www.lintcode.com/problem/k-diff-pairs-in-an-array/description

import java.util.Map;
import java.util.TreeMap;

public class Solution {
    /**
     * @param nums: an array of integers
     * @param k: an integer
     * @return: the number of unique k-diff pairs
     */
    public int findPairs(int[] nums, int k) {
        // Write your code here
        TreeMap treeMap = new TreeMap<>();
        for (int i = 0; i < nums.length; i++) {
            int num = nums[i];
            Integer count = treeMap.get(num);
            if (count == null) {
                count = 1;
            } else {
                count++;
            }
            treeMap.put(num, count);
        }
        int res = 0;
        while (!treeMap.isEmpty()) {
            Map.Entry firstEntry = treeMap.pollFirstEntry();
            Integer key = firstEntry.getKey();
            int target = key + k;
            if (target != key) {
                if (treeMap.containsKey(target)) {
                    res++;
                }
            } else {
                if (firstEntry.getValue() >= 2) {
                    res++;
                }
            }
        }
        return res;
    }
}