Squares of a Sorted Array

See the original problem on HackerRank.

Given an array of integers A sorted in non-decreasing order, return an array of the squares of each number, also in sorted non-decreasing order, in linear time.

Input Format

The first line contains N, the length of the array. The second line contains N space-separated integers, the elements of the array.

Constraints

  • \(1 \leq N \leq 10^6\)
  • \(-10^3 \leq A[i] \leq 10^3\)

Output Format

Print the array of the squares of each number, also in sorted non-decreasing 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 in-place:

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 non-negative
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, " "));
We've worked on this challenge in these gyms: modena 
comments powered by Disqus