/******************************************************************************
Welcome to GDB Online.
GDB online is an online compiler and debugger tool for C, C++, Python, Java, PHP, Ruby, Perl,
C#, VB, Swift, Pascal, Fortran, Haskell, Objective-C, Assembly, HTML, CSS, JS, SQLite, Prolog.
Code, Compile, Run and Debug online from anywhere in world.
*******************************************************************************/
import java.util.Arrays;
import java.util.Comparator;
import java.util.concurrent.ThreadLocalRandom;
public class Main
{
static class Elem implements Comparable<Elem>{
int value;
int index;
long dist;
public Elem(int value, int index){
this.value = value;
this.index = index;
}
@Override
public int compareTo(Elem o){
return Integer.compare(value, o.value);
}
}
public static void main(String[] args){
int n = 300_000;
int k = 3_000;
Elem[] elems = new Elem[n];
for(int i = 0; i < n; i++){
elems[i] = new Elem(ThreadLocalRandom.current().nextInt(), i);
}
solve(n, k, elems);
}
private static void solve(int n, int k, Elem[] elems){
Arrays.sort(elems);
long time = System.nanoTime();
for(int i = 0; i < n; i++){
elems[i].dist = findDistForIth(elems, i, k);
}
// Arrays.sort(elems, Comparator.comparingInt(elem -> elem.index));
// System.out.print(elems[0].dist);
// for(int i = 1; i < n; i++){
// System.out.print(" " + elems[i].dist);
// }
System.out.println((System.nanoTime() - time)/1_000_000_000.0);
}
private static long findDistForIth(Elem[] elems, int i, int k){
int midElem = elems[i].value;
int left = i - 1;
int right = i + 1;
long dist = 0;
for(int j = 0; j < k; j++){
if(left < 0){
dist += elems[right++].value - midElem;
}else if(right >= elems.length){
dist += midElem - elems[left--].value;
}else{
int leftAdd = midElem - elems[left].value;
int rightAdd = elems[right].value - midElem;
if(leftAdd < rightAdd){
dist+=leftAdd;
left--;
}else{
dist+=rightAdd;
right++;
}
}
}
return dist;
}
}
