See the original problem on HackerRank.
Given an array of integers A
sorted in nondecreasing order, return an array of the squares of each number, also in sorted nondecreasing order, in linear time.
The first line contains N
, the length of the array.
The second line contains N
spaceseparated integers, the elements of the array.
Constraints
 \(1 \leq N \leq 10^6\)
 \(10^3 \leq A[i] \leq 10^3\)
Print the array of the squares of each number, also in sorted nondecreasing order.
Solutions
The naive solution consists in creating an array of the squares of each element, and sort them:
1
2
3
4
5
6
7
8
9
10
11
12

class Solution
{
public int[] sortedSquares(int[] A) {
int N = A.length;
int[] ans = new int[N];
for (int i = 0; i < N; ++i)
ans[i] = A[i] * A[i];
Arrays.sort(ans);
return ans;
}
}

A variation consists in mapping the squared vector into a sorted data structure and map it back to a vector:
1
2
3
4
5
6

vector<int> sortedSquares(vector<int>& A)
{
multiset<int> s;
transform(begin(A), end(A), inserter(s, begin(s)), [](int i){ return i*i; });
return vector<int>(s.begin(), s.end());
}

Clearly, such solutions are both order of \(O(N \cdot LogN)\). The latter is even worse since is uses more space.
We asked to find a linear time algorithm.
A more clever solution comes with an intuition: the array might have a negative part and a following positive part. We can iterate both at the same time, adding to the result array the lowest value for each iteration.
First of all, we need to find the beginning of the positive values. This is a simple array find. Then, we can set two indexes: one moving backwards from the last negative value, and one moving forward from the first positive value.
At every step, we choose which element has the lowest magnitude and add its square to the output array (we print it). We move one of the two indexes accordingly.
Here is a C++ implementation:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

vector<int> sortedSquares(vector<int>& A)
{
vector<int> res(A.size());
// locate the first positive value
auto it = find_if(begin(A), end(A), [](int i) { return i >= 0; });
auto pos = (int)distance(begin(A), it);
auto neg = pos  1;
int len = (int)A.size();
for (auto i = 0; i < len; ++i)
{
// negative values finished
if (neg < 0)
{
res[i] = square(A[pos++]);
}
// positive values finished
else if (pos >= len)
{
res[i] = square(A[neg]);
}
// positive value is bigger
else if (square(A[pos]) > square(A[neg]))
{
res[i] = square(A[neg]);
}
// negative value is bigger
else
{
res[i] = square(A[pos++]);
}
}
return res;
}

An alternative solution runs in a single pass and still consists in using two indexes but we start filling the result array backwards this time.
We start comparing the absolute value of the first and the last element:
 if the leftmost is bigger, then we add the square of the rightmost and move the right index back by one
 otherwise, we add the square of the leftmost and move the left index up by one
Here is the implementation:
1
2
3
4
5
6
7
8
9
10
11
12
13

vector<int> sortedSquares(vector<int>& A)
{
vector<int> ret(A.size());
int r=A.size()1,l=0;
for(int i=r;i>=0;i)
{
if(abs(A[r])>abs(A[l]))
ret[i]=A[r]*A[r];
else
ret[i]=A[l]*A[l++];
}
return ret;
}

Another approach using two lists by Eduard “Eddie” Rubio Cholbi and Stefano Ligabue:
1
2
3
4
5
6
7
8

def squaresOfSortedArray(values):
positive = []
negative = []
for v in values:
(negative, positive)[v >= 0].append(v**2)
return merge(reversed(negative), positive)

The idea is to fill two lists, one with the square of positive values and one with the square of negative values, and then merge them together (since they are sorted).
Note the pearl of wisdom: using the selector inside the array access operator to select the former or the latter array depending on the sign of the current value.
After seeing this soluton at the retrospective in Modena (April 2019), Marco Arena had an insight and he wrote a very similar solution in C++ that is totally inplace:
1
2
3
4
5
6
7

int N; cin >> N;
vector<int> v(N);
copy_n(istream_iterator<int>(cin), N, begin(v));
auto firstPos = find_if(begin(v), end(v), [](auto i){ return i>=0; }); // find first nonnegative
transform(begin(v), end(v), begin(v), [](auto i){return i*i;}); // square them all
merge(make_reverse_iterator(firstPos), rend(v), firstPos, end(v), ostream_iterator<int>(cout, " ")); // the magic is here

This snippet is a tribute to the flexibility of the C++ standard library.
Here is a possible application of C++20 ranges to the solution above:
1
2
3
4
5
6
7
8

auto firstPos = ranges::find_if(v, [](auto i){ return i>=0; });
auto negatives = views::counted(std::make_reverse_iterator(firstPos), std::distance(std::begin(v), firstPos));
auto positives = views::counted(firstPos, std::distance(firstPos, std::end(v)));
const auto square = [](auto i) { return i*i; };
ranges::merge(views::transform(positives, square),
views::transform(negatives, square),
ranges::ostream_iterator(std::cout, " "));

Some more exploration of the library gets us to this more succinct way of creating the source ranges:
1
2
3
4
5
6
7
8

auto firstPos = ranges::find_if(v, [](auto i){ return i>=0; });
auto positives = ranges::subrange(firstPos, std::end(v));
auto negatives = ranges::subrange(std::make_reverse_iterator(firstPos), std::rend(v));
const auto square = [](auto i) { return i*i; };
ranges::merge(views::transform(positives, square),
views::transform(negatives, square),
ranges::ostream_iterator(std::cout, " "));

Here is another version that uses drop_while
:
1
2
3
4
5
6
7

auto positives = views::drop_while(v, [](auto i){ return i<0; });
auto negatives = views::drop_while(views::reverse(v), [](auto i) { return i>=0; });
const auto square = [](auto i) { return i*i; };
ranges::merge(views::transform(positives, square),
views::transform(negatives, square),
ranges::ostream_iterator(std::cout, " "));

This one is actually a bit less efficient because the “partition point” is found twice.
Another approach mimics the previous Python solution but it’s clearly less efficient and more verbose:
1
2
3
4
5
6
7
8
9
10

int N; cin >> N;
vector<int> v(N);
copy_n(istream_iterator<int>(cin), N, begin(v));
vector<int> negatives, positives;
partition_copy(begin(v), end(v), back_inserter(negatives), back_inserter(positives), [](auto i) { return i<0; });
reverse(begin(negatives), end(negatives));
transform(begin(negatives), end(negatives), begin(negatives), [](auto i){ return i*i; });
transform(begin(positives), end(positives), begin(positives), [](auto i){ return i*i; });
merge(begin(negatives), end(negatives), begin(positives), end(positives), ostream_iterator<int>(cout, " "));
