My first try was, using partial sort to figure out numbers layer by layer in the heap.. it only failed with TLE with the last test case. The problem is, partial sort cannot guaratee O(n) every time.
class Solution
{
void kth(vector<int> &A, int s, int e, int k) // all zero based
{
if(s >= e) return; // Partition
int i = e, j = e;
{
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<> dis(s, e); int pi = dis(gen);
int pivot = A[pi];
swap(A[pi], A[s]); while(j > s)
{
if(A[j] >= pivot)
{
swap(A[i], A[j]);
i --; j = i;
}
else
{
j --;
}
}
swap(A[i], A[s]);
} // Recursion
if(i < k)
{
kth(A, i + , e, k);
}
else if(i > k)
{
kth(A, s, i - , k);
}
}
public:
/**
* @param A: Given an integer array
* @return: void
*/
void heapify(vector<int> &A)
{
size_t n = A.size();
int s = , e = n - ;
while (s < e)
{
int cnt = e - s + ;
int h = ceil(log2(cnt));
int k = (pow(, h) - )/;
kth(A, s, e, k - );
e = k - ;
}
}
};
A smarter way is as below. Its strategy is "per-node maintanence".
class Solution {
void help(vector<int> &A, int i)
{
int n = A.size();
int li = i * + , ri = i * + ;
int left = li < n ? A[li] : INT_MAX;
int right= ri < n ? A[ri] : INT_MAX;
if(left < right && left < A[i])
{
swap(A[li], A[i]);
help(A, li);
}
else if(right < left && right < A[i])
{
swap(A[ri], A[i]);
help(A, ri);
}
}
public:
/**
* @param A: Given an integer array
* @return: void
*/
void heapify(vector<int> &A) {
for(int i = A.size() / ; i >= ; i --)
help(A, i);
}
};