What's time complexity of following modified bucket sort solution

It's kind of bucket sort algorithm, trying to get K nearest locations from point (0,0). This is being done by calculating distances of these locations and bucketing them based on distance. If two locations are at equal distance then priority given to location with closet x and then y(if x values are same)

What is time complexity of following solution? O(NlogN) or O(KNlogN) or anything else

// java

Input: int k, List<Location>
Output: List<Location>

public class Location {
    //constructor x and y
    int x;
    int y;
    //get 
    //set

    //hashcode and equals
}

public class Solution {
    public List<Location> returnKNearestLocation(int k, List<Location> locations) {
        Map<Float, Set<Location>> map = new TreeMap<>();
        for(Location loc: locations) {
            float dist = calculateDistance(new Location(0,0), loc);
            List<Location> temp = map.getOrDefault(dist, new HashSet());
            temp.add(loc);
            map.put(temp);
        }

        List<Location> result = new ArrayList<>();
        int size = k;
        while(size > 0) {
            for(Float key : map.keySet()) {
                Set<Location> loc = map.get(key);
                Collection.sort(loc, p1, p2 -> {
                        return p1.x.equals(p2.x)? p1.y.compare(p2.y) : p1.x.compare(p2.x);
                });

                for(Location currLoc : loc) {
                    result.add(currLoc);
                    if(result.size() == k) {
                        return result;
                    }
                }
                size = size - loc.size();
            }                
        }

        return result;
    }

    private float calculateDistance(Location p, Location q) {
       // calculate distance Math.sqrt((P.x - Q.x)^2 + (P.y - Q.y)^2)
    }
}

I'd say O(N*log(N)).

• Assuming float dist = calculateDistance(new Location(0,0), loc) is O(1).
• Implementations of put() and getOrDefault() in a TreeMap are O(log(N)). Other implementations offer O(1). I'd take a look at HashMap.
• Implementation of add in a HashSet is O(1) (amortized).
• Implementation of add in an ArrayList is O(1) (amortized)
• BucketSort has a worst-case of O(N^2), in the case that all items are placed in the same bucket.

I'd take a look at HashMap instead of TreeMap.

So, assuming I'm not wrong:

• The first loop (foreach loc) has a complexity of O(N*log(N)).
• The second loop (while size > 0) is O(N*log(N)).
  • Assuming all locations are placed in the same bucket, we'll sort all the items (O(N*log(N))).
  • Then we'll iterate through the first K elements, and add to result -- this takes O(K) < O(N).