Encuentra los K elementos mas cercanos.

Ok, una de las cosas que me distrae y me gusta es hacer estos problemas de algoritmia que puede que te salgan en alguna entrevista, para esto utilizo LeetCode para encontrar ejercicios y practicar.

Entonces, esta vez voy a plantear mi razonamiento y solución a un problema de LeetCode, el problema "658 - Find the K closest elements". El ejercicio es el siguiente:

Given a sorted integer array arr, two integers k and x, 
return the k closest integers to x in the array. 
The result should also be sorted in ascending order.
 
  An integer a is closer to x than an integer b if:
 
  |a - x| < |b - x|, or
  |a - x| == |b - x| and a < b
 
 
  Example 1:
  Input: arr = [1,2,3,4,5], k = 4, x = 3
  Output: [1,2,3,4]
  
  Example 2:
  Input: arr = [1,2,3,4,5], k = 4, x = -1
  Output: [1,2,3,4]
 
  Constraints:
  1 <= k <= arr.length
  1 <= arr.length <= 104
  arr is sorted in ascending order.
  -104 <= arr[i], x <= 104 

El problema nos pide encontrar los "K" números enteros más cercanos al valor "X" dentro del array y formar un array con esos elementos. Vamos a analizar el primer ejemplo para entender bien el ejercicio.

Input: arr = [1,2,3,4,5], k = 4, x = 3 ()
Output: [1,2,3,4]

arr => es el arreglo donde vamos a búscar los elementos que tenemos que retornar.
K => cuantos elementos tenemos que retornar.
3 => valor de referencia para hacer la búsqueda. Vamos a llamarlo "X"

La primera pregunta que me haría es por qué el Output: [1,2,3,4]? Vamos a ver cómo llegamos a ese resultado. Como dice el ejercicio un número "a" está más cerca de "X" que otro número "b" si:

  |a - x| < |b - x|, or
  |a - x| == |b - x| and a < b

Si vemos el primer ejemplo y comparamos sus elementos con el valor de referecia siguiendo las reglas de arriba, entonces tenemos:

Valor de referencia X = 3.
Input = [1,2,3,4,5]
Output = [x,x,x,x] // tendrá esta forma por K = 4.

|1-3| = 2
|2-3| = 1
|3-3| = 0
|4-3| = 1
|5-3| = 2

Ahora elegimos los 4 elementos más cercanos, que serían los que tienen el menor resultado. Viendo uno a uno los resultados el menor es el 3, ya tenemos un espacio del arreglo output ocupado y nos faltan tres. Luego le siguen el 2 y el 4 que dieron como resultado 1, en este caso elegimos los dos porque caben dentro del output y con esto ya tenemos 3 elementos en el output. Nos falta el último elemento del output y nos quedan dos elementos por comparar en el input, el 1 y el 5, ambos con la misma distancia al valor de referencia, |1-3| = |5-3| = 2, es decir un empate y los empates se resuelven tomando el elemento menor del input, en este caso el 1, entonces nuestro output es:

Output [1,2,3,4,5]

Excelente, ya entendmos el problema, su input, su output y las reglas necesarias para hacer la comparación. Vamos ahora a plantear un algoritmo para resolver el problema, capaz algo como lo siguiente funcione:

1. Tomar cada elemento del arreglo.
2. Calcular la diferencia de cada elemento con X y guardarlo en un arreglo junto con el valor del elemento.
3. Realizar un ordenamiento estable de ese arreglo según la diferencia calculada.
4. Retornar los K primeros elementos.

Bueno, para ser honesto esto se me ocurrió mientras escribía y no sé si funciona del todo, prefiero ir con mi solución que ya la codie y la probé. Teniendo en cuenta que el arreglo input está ordenado podemos plantear un algoritmo como el siguiente:

1. Buscar el índice del valor de referencia en el arreglo de entrada (Binary search).
   • Si el valor de referencia está dentro del arreglo voy a tener su índice dentro del arreglo.
   • Si el valor de referencia NO está dentro del arreglo voy a tener un índice dentro del arreglo para comparar.
   • Como el arreglo está ordenado, los elementos más cercanos van a estar al rededor de ese índice.
2. Coloco un puntero en el índice que encontré en el paso 1, llamemoslo (right) y otro puntero justo despues, llamemoslo (left).
3. Comparo right y left y me quedo con el que esté más cerca según la regla y lo agrego al arreglo result.
4. Muevo right o left dependiendo de que valor seleccione en el paso 3, así:
   • Si en el paso 3 seleccione right, entonces muevo una posición left.
   • Si en el paso 3 seleccione left, entonces muevo una posición right.
5. Repito los pasos 3 y 4 hasta completar el arreglo.
6. Devuelvo el arreglo ordenado.

Esta solución si estoy seguro de que funciona jajaja... bueno, ahora veamos el código.

1. Buscar el índice del valor de referencia en el arreglo de entrada (Binary search).
   • Si el valor de referencia está dentro del arreglo voy a tener su índice dentro del arreglo.
   • Si el valor de referencia NO está dentro del arreglo voy a tener un índice dentro del arreglo para comparar.
   • Como el arreglo está ordenado, los elementos más cercanos van a estar al rededor de ese índice.
2. Coloco un puntero en el índice que encontré en el paso 1, llamemoslo (right) y otro puntero justo despues, llamemoslo (left).

var result = new ArrayList<Integer>(); // Inicializo el arreglo resultado.
var right = binarySearch(arr, 0, arr.length, x); // puntero right
var left = right - 1; // puntero left

3. Comparo right y left y me quedo con el que esté más cerca según la regla y lo agrego al arreglo result.
4. Muevo right o left dependiendo de que valor seleccione en el paso 3, así:
   • Si en el paso 3 seleccione right, entonces muevo una posición left.
   • Si en el paso 3 seleccione left, entonces muevo una posición right.

while(result.size() < k){  
    var rightInBound = right >= 0 && right < arr.length;  
    var leftInBound = left >= 0 && left < arr.length;  
  
    if(!leftInBound && rightInBound) {//si no hay elementos en left, tomo right.
        result.add(arr[right]);  
        right++;  
    }else if(!rightInBound && leftInBound) {//si no hay elementos right, tomo left.
        result.add(arr[left]);  
        left--;  
    } else {  //Comparo los elementos y muevo los punteros.
        var mod = Math.abs(arr[left] - x) - Math.abs(arr[right] - x);  
        if(mod <= 0) {  
            result.add(arr[left]);  
            left--;  
        } else if(mod > 0) {  
            result.add(arr[right]);  
            right++;  
        }  
    }  
  
}

6. Devuelvo el arreglo result ordenado.

return result.sort(Integer::compareTo);  

Podemos usar LeetCode para ver los casos de prueba.

Las corridas fallidas fueron por un error en el binarySearch. Falta algo, time complexity, el binarySearch es O(log(n)) asi que es bastante rápido, más la comparación sería O(k) porque la iteración la hacemos K veces, es decir, es lineal, y habría que agregarle la complejidad del algoritmo de ordenamiento del final; en resumen sería:

O(log(arr.length)) + O(k) + O(ordenamiento)

Pueden ver el código completo en mi repo de GitHubnto que encontré (right) y otro a su lado (left).