Find all k points which are closest to origin We have a list of points on the plane.  Find the K closest points to
  the origin (0, 0).
  
  (Here, the distance between two points on a plane is the Euclidean
  distance.)
  
  You may return the answer in any order.  The answer is guaranteed to
  be unique (except for the order that it is in.)
  
  Example 1:
  
  Input: points = [[1,3],[-2,2]], K = 1
  
  Output: [[-2,2]]
  
  Explanation: The distance between (1, 3) and the origin is
  \$\sqrt{10}\$.  The distance between (-2, 2) and the origin is
  \$\sqrt{8}\$.  Since \$\sqrt{8} &lt; \sqrt{10}\$, (-2, 2) is closer to
  the origin.  We only want the closest K = 1 points from the origin, so
  the answer is just [[-2,2]].
  
  Example 2:
  
  Input: points = [[3,3],[5,-1],[-2,4]], K = 2
  
  Output: [[3,3],[-2,4]]
  
  (The answer [[-2,4],[3,3]] would also be accepted.)
  
  Note:
  
  
  1 &lt;= K &lt;= points.length &lt;= 10000
  -10000 &lt; points[i][0] &lt; 10000
  -10000 &lt; points[i][1] &lt; 10000
  
  
  Also note that there can be a situation where distance of 2 nodes are
  equal and hence can print any of the node.




Here is my approach.


Took tree map (So that I get all sorted distance)
Start filling tree map for all points
Finally took first k element.




public class KClosestPointsToOrigin {

    public int[][] kClosest(int[][] points, int K) {

        int rows = points.length;


        SortedMap&lt;Double, List&lt;CoordinatePoint&gt;&gt; distanceToPoint = new TreeMap&lt;&gt;();

        for(int i =0; i&lt;rows; i++) {
            double distance = getDistance(points[i]);
            if(Objects.nonNull(distanceToPoint.get(distance))) {
                List&lt;CoordinatePoint&gt; coordinatePointList = distanceToPoint.get(distance);
                coordinatePointList.add(new CoordinatePoint(points[i][0], points[i][1]));
            } else {
                List&lt;CoordinatePoint&gt; coordinatePoints = new ArrayList&lt;&gt;();
                coordinatePoints.add(new CoordinatePoint(points[i][0], points[i][1]));
                distanceToPoint.put(distance, coordinatePoints);// x and y coordinates.
            }
        }

        int[][] arrayToReturn = new int[K][2];
        int counter = 0;
        for (Double key : distanceToPoint.keySet()) {
            List&lt;CoordinatePoint&gt; coordinatePoints = distanceToPoint.get(key);
            Iterator iterator1 = coordinatePoints.iterator();
            while (iterator1.hasNext() &amp;&amp; counter &lt; K) {
                CoordinatePoint coordinatePoint = (CoordinatePoint) iterator1.next();
                arrayToReturn[counter][0] = coordinatePoint.x;
                arrayToReturn[counter][1] = coordinatePoint.y;
                counter++;
            }
        }

        return arrayToReturn;
    }

    private double getDistance(int[] point) {
        int x = point[0];
        int y = point[1];

        int x2 = Math.abs(x) * Math.abs(x);
        int y2 = Math.abs(y) * Math.abs(y);

        return Math.sqrt(x2+y2);
    }
}

class CoordinatePoint {
    int x;
    int y;

    public CoordinatePoint(int x, int y) {
        this.x = x;
        this.y = y;
    }
}


However, this solution is not efficient as runtime and memory usage is high. Can you please help me to optimize the solution.
This task sounds as if it came directly from an advertisement for the Java 8 streams API:


Sort the points by their distance
Take the minimum k points


This boils down to the following code:

static int[][] kClosest(int[][] points, int k) {
    return Arrays.stream(points)
        .sorted(Comparator.comparing((int[] point) -&gt; point[0] * point[0] + point[1] * point[1]))
        .limit(k)
        .toArray(int[][]::new);
}


That code is written from the top of my head. Since the Java streams API is general-purpose and intended to be highly optimized, it should notice that it only ever needs to remember the \$k\$ smallest elements. You should check this by counting how often the distance function is called. To do that you should extract it to a local method, which is something that your IDE can do for you.

Compared to your code:


I left out the Math.sqrt since the point coordinates are restricted to 10000, which means that the square of the distance can become at most 2e8, which luckily is a bit smaller than Integer.MAX_VALUE.
Calling Math.abs is unnecessary as well.
Making a map of lists may help in pathological cases when the given points all have the same distance. In other cases it can be left out.