Official Editorial (C++)

Explanation

There's a ton of ways to solve this, but since you're probably here from the square root decomposition module, let's just go over that one. The official editorial does have some nice alternative solutions as well.

Subtask

Let's first try to solve the subtask and come up with an $\mathcal{O}(N^2)$ solution.

Let's try to construct the minimum possible ordering for each prefix of the array. We can do this with a strategy similar to insertion sort.

We can insert an element into the prefix by swapping it down to the left until it reaches position $i$ where $|a[i] - a[i - 1]| > k$ because its impossible to keep swapping it. However, if $a[i] > a[i + 1]$, it is more optimal to swap it back to the right until it reaches a point where $a[i] < a[i + 1]$.

Full solution

We can speed this up with square root decomposition. We can split the array into blocks of size $\sqrt{n}$.

We know an element $a[i]$ can pass completely through a block if $a[i] + k \ge \texttt{MAX}$ and $a[i] - k \le \texttt{MIN}$ where $\texttt{MAX}$ and $\texttt{MIN}$ are the maximum and minimum elements of the block respectively. We can find the rightmost block that element $a[i]$ cannot pass through which takes at most O($\sqrt{n}$) since we check every block.

We can insert this element into the square root block at its correct position by swapping it down to the left inside the block till it reaches a place where it can't be swapped down anymore. This also takes $\mathcal{O}(\sqrt{N})$ time because of the block size.

Now we need to swap it back towards the right until it reaches a point where which makes the array lexographically minimum. We can do this by finding the first block to the right of it that has it's $\texttt{MAX} > a[i]$ by checking each block again in $\mathcal{O}(\sqrt{N})$.

Notice that there is an edge case here where $a[i]$ can never be swapped down to the left past another element that is greater than it because of the restriction of $k$ even thought the $\texttt{MAX}$ element is greater than $a[i]$. In this case, we can push $a[i]$ into the next block and then apply the strategy of checking when the $\texttt{MAX} > a[i]$. We can check this by seeing that if after inserting $a[i]$ into the block and swapping towards the left and right appropriately as described above, $a[i]$ reaches the end of the block.

Now the problem arises that when we add enough elements, the size of a block could potentially be bigger than $\sqrt{n}$. To deal with this issue we can split the block into two smaller blocks (each half the size of the original block) and insert them into the list of blocks. Inserting into a vector in C++ or ArrayList in Java, is dependent on the size of the array which in our case which takes $\mathcal{O}(\sqrt{N})$. We can prove that the amount of blocks is bounded to $2\sqrt{n}$.

Processing each element of the array took $\mathcal{O}(\sqrt{N})$ time so the total time complexity is $\mathcal{O}(N\sqrt{N})$.

Implementation

Time Complexity: $\mathcal{O}(N\sqrt{N})$

Copy
#include <bits/stdc++.h>
using namespace std;

int n, k;

struct Block {
	int hi = -INT32_MAX;
	int lo = INT32_MAX;
	vector<int> vals;

Copy
import java.io.*;
import java.util.*;

public class MinimizingHaybales {
	// sqrt(1e5) is roughly 300
	static final int BLOCK_SIZE = 300;
	static int k;

	public static class Block {
		public int hi, lo;