package questions;

import java.util.*;

public class FindKClosestToCenter  {
	
	//Solution 1, Array sorting, Time worst O(n^2), average O(n), Space O(n), n is number of points
	public static int[][] kClosest(int[][] points, int k) {
		Arrays.sort(points, (a, b) -> compare(a, b));
		return Arrays.copyOfRange(points, 0, k);
	}
	
	//Solution 2, quick select, Time worst O(n^2) average O(n), Space worst O(n) average O(logn)
	public static int[][] kClosestQuick(int[][] points, int k) {
	     int n =  points.length, low = 0, high = n - 1;
	     while (low <= high) {
	    	 int mid = partition(points, low, high);
	    	 if (mid == k) 
	    		 break;
	    	 if (mid < k) 
	        	low = mid + 1; 
	    	 else 
	            high = mid - 1;
	    }
	    return Arrays.copyOfRange(points, 0, k);
	}

	//Partition, two pointers, Time worst O(n^2) average O(n), Space O(1)
	private static int partition(int[][] points, int low, int high) {
	    int[] pivot = points[low];
	    while (low < high) {
	        while (low < high && compare(points[high], pivot) >= 0) {
	        	high--;
	        }
	        points[low] = points[high];	        
	        while (low < high && compare(points[low], pivot) <= 0){ 
	        	low++;
	        }
	        points[high] = points[low];
	    }    
	    points[low] = pivot;	    
	    return high;
	}
	
	//Solution 3, PriorityQueue, Time O(nlogn), space O(n)
	public static int[][] kClosestPQ(int[][] points, int k) {
	    PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> compare(a, b)); //sort by distance, ascending 
	    for (int[] p : points) 
	        pq.add(p);    
	    int[][] res = new int[k][2];
	    for (int i = 0; i < k; i++)
	    	res[i] = pq.poll();
	    return res;
	}
	
	//Compare based on Euclidean distance formula, Time O(1), Space O(1)
	private static int compare(int[] a, int[] b) {
	    return (a[0]*a[0] + a[1]*a[1]) - (b[0]*b[0] + b[1]*b[1]);
	}
			
	//Utility print points, Time O(n), Space O(1)
	private static void printPoints(int[][] points) {
		for (int i = 0; i < points.length; i++) 
			System.out.print("(" + points[i][0] + ", " + points[i][1] + ") ");
		System.out.println();
	}
	 
	public static void main(String[] args) throws Exception { 
		
		int[][] points = { {4,5}, {1,3}, {3,2},{2, 4}, {2, 3},{3,7}}; 
        int k = 2;
        
		//Solution 1, Array sorting
		int[][] res = kClosest(points, k);
		System.out.println("sorting: ");
		printPoints(res);
        
		//Solution 2, quick select
		int[][] res1 = kClosestQuick(points, k);
		System.out.println("quick select: ");
		printPoints(res1);
		
		//Solution 3, Priority Queue
		int[][] res2 = kClosestPQ(points, k);
		System.out.println("PriorityQueue: ");
		printPoints(res2);
    }
}