LeetCode 977: Squares of a sorted array, maintaining non-decreasing order

Question

I tried solving the LeetCode question like many others, trying to incorporate the O(n) time complexity requirement.

977. Squares of a Sorted Array

Given an integer array nums sorted in non-decreasing order, return an array of the squares of each number sorted in non-decreasing order.

Example

• Input: nums = [-4,-1,0,3,10]
• Output: [0,1,9,16,100]
• Explanation: After squaring, the array becomes [16,1,0,9,100]. After sorting, it becomes [0,1,9,16,100].

I came up with the following:

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        n = len(nums)
        i = 0
        j = n -1
        # binary search O(log n)
        if nums[0] == 0 or nums[0]> 0:
            for i in range(n):
                nums[i] = nums[i]**2
            return nums
        elif nums[n-1] == 0 or nums[n-1] < 0:
            nums.reverse()
            for i in range(n):
                nums[i] = nums[i]**2
            return nums
        k1 = 0
        k = n-1
        while i<j and k!=k1:
            k1 = k
            k = (i+j)//2  
            if nums[k] == 0:
                break
            elif nums[k] > 0:
                j = k
            else:
                i = k
        if abs(nums[k]) > abs(nums[k+1]):
            k = k+1
        # sorting this way should cost O(n) time
        if k < (n-1)//2:
            q = k+1
            for s in range(k-1,-1,-1):
                p = nums[s]
                for l in range(s,q-1):
                    nums[l] = nums[l+1]
                while q < n and abs(p) > abs(nums[q]):
                    nums[q-1] = nums[q]
                    q += 1
                nums[q-1] = p
        else:
            q = k-1
            for s in range(k+1,n):
                p = nums[s]
                for l in range(s-1,q,-1):
                    nums[l+1] = nums[l]
                while q >= 0 and abs(p) > abs(nums[q]):
                    nums[q+1] = nums[q]
                    q -= 1
                nums[q+1] = p   
            nums.reverse()
        # this for-loop also costs O(n)
        for i in range(n):
            nums[i] = nums[i]**2
        # therefore O(log n)+O(n)+O(n) -> O(n)
        return nums

by manually calculating the time cost (and also the AI tool by LeetCode). I seem to get O(n) time complexity, but this code is incredibly slow. Is it due to the 'chunkiness' of it or something else? Feel free to point out any redundancies or inefficient lines.

Answer 1

modern annotations

These are lovely, thank you.

    def sortedSquares(self, nums: List[int]) -> List[int]:

They're very clear. But for modern cPython interpreters, 3.9+, prefer list[int] without the from typing import List.

Also, PEP 8 asks that you spell it sorted_squares() (unless LeetCode imposes some weird non-pythonic requirement on folks offering python submissions).

comments lie!

They lie all too often. People don't mean for it to happen; it just does.

        # binary search O(log n)
        if nums[0] == 0 or nums[0]> 0:
            for i in range(n):
                nums[i] = nums[i]**2
            return nums
        elif nums[n-1] == 0 or nums[n-1] < 0:
            nums.reverse()
            for i in range(n):
                nums[i] = nums[i]**2
            return nums

This in no way resembles a binary search. The relevant remark would explain that we're seeing if we're in the special case of "all positive" or "all negative".

>=

        if nums[0] == 0 or nums[0]> 0:

That's an odd way to phrase if nums[0] >= 0:

        elif nums[n-1] ...

That's perfectly fine, but the idiomatic way to talk about the last element would be nums[-1].

In any event, you accomplished the special casing in \$O(n)\$ linear time, so no harm done. Mostly we're trying to avoid the extra \$\log n\$ factor that sorting would introduce.

meaningful identifiers

        k1 = 0
        k = n-1

The i and j indexes were great, but these two are terrible terrible identifiers. Pick names that show what they mean, or at least offer a # comment to help out the Gentle Reader. The and k!=k1 conjunct suggests that you have some loop invariant in mind, but you didn't go to the trouble of telling us what it is, how it relates to the business problem at hand.

The k = (i+j)//2 expression suggests that mid or midpoint might be a relevant name to use.

            q = k+1

I don't know what this means.

If you were hoping to communicate an idea to the machine, good, chalk up a win. If you were hoping to communicate a technical idea to a collaborator, this attempt was not as effective as it could have been. Use English narrative, or descriptive names, or at least the crutch of # comments.

extract helper

            if nums[k] == 0:
                break
            elif nums[k] > 0:
                j = k

Why are these even special cases?

The real thing you wish to do is find a sign change. Recall that \$O(n + \log n) = O(n)\$. So, bonus points for the binary search, I guess? But a simple linear search for sign change would have remained within your \$O(n)\$ budget.

The if k < (n-1)//2: loop offers another fruitful opportunity to Extract Helper. Minimally you get to def foo() where "foo" usefully describes what's happening. Ideally you would add a """docstring""" that offers further details.

Recall that the bisect module stands at the ready to help with many of your binary search needs.

---

algorithm

There's only two cases here. A number is negative or it isn't. I would expect a solution to follow 2-way merge:

1. Linear scan to find change-of-sign (if any).
2. Maintain indices i and j which go outward from zero. While absolute value of one or the other is "winning", emit its square. Then keep iterating.

Both steps have constant memory cost and linear time complexity.

Given that we already know the size of the output array, we don't even need to scan for where the sign change happens, since we can fill it in in reverse. Simply init i & j to start and end, and do 2-way merge inward, emitting squares as you go, in reverse order. At some point the two indexes will locate where the sign change is, and then we're done.

constant time solution

Notice the constraints.

There are at most 10_000 elements, and their absolute values shall be at most 10_000. It's not hard to allocate a container of that size. Simply store absolute values of the inputs, and then read them out in sequence. No \$\log n\$ term involved at all. Is this cheating? Yes. Kind of. The key is to figure out the rules, and play within them. (And benchmark against the previous approach. Which this won't be competitive with, for most input sizes.)

Answer 2

I will try not repeat comments/suggestions made by others.

Do Not Modify the Input

As I interpret the question, the input array (list) should be left unmodified. Your solution clearly is mutating the input.

Encapsulation

Whereas it is perfectly okay to encapsulate your solution within a class, there is no reason for method sortedSquares (which I would rename to sorted_squares in line with the PEP8 style guide), not to be a static or class method.

Is the Algorithm Truly O(n)?

I admit to having fogged over a bit trying to figure out what is being done (the absence of comments did not help). One of the few comments states with the wave of the hand that the time complexity is O(n), but it is not apparent to me because of the code's complexity that this is the case. You have a comment that states # sorting this way should cost O(n) time. But looking at the code that follows, you seem to be comparing values of the input. Any algorithm that includes sorting where the sorting algorithm is based on the comparison of input values will be at best O(n * log(n)).

A Linear Algorithm

Here the basic idea is to loop through each element of the list to find the index, idx, of the first element that is non-negative. This is O(n). There are 3 cases:

1. idx is 0. This is a special case where all elements are non-negative.
2. idx is equal to n_elements, the length of the list. This is a special case where all elements are negative.
3. Otherwise, we perform a merge between two subarrays of the list.

Note that in the case where we are doing a merge, the time complexity is O(m + n) where m and n are the lengths of the two subarrays. But m + n == n_elements, so the time complexity is just O(n_elements).

Note also that the following code uses a sentinel value of math.inf to simplify the logic.

def solution(arr: list[int|float]) -> list[int|float]:
    """Given a list arr where successive elements are
    non-decreasing, return a list whose values are the squares
    of the elements in non-decreasing order."""

    from math import inf

    n_elements = len(arr)
    idx = 0
    while idx < n_elements and arr[idx] < 0:
        idx += 1

    # Two special cases:
    if idx == 0:
        # All non-negative:
        return [x * x for x in arr]
    if idx == n_elements:
        # All negative:
        return [x * x for x in arr[::-1]]

    # idx is the index of the first non-negative element
    # Create generator expressions for processing each subarray:
    arr1 = (arr[i] ** 2 for i in range(idx - 1, -1, -1))
    arr2 = (arr[i] ** 2 for i in range(idx, n_elements))
    # Allocate the results
    results = [None] * n_elements

    n1 = idx  # Size of the first subarray
    n2 = n_elements - n1  # Size of the second subarray
    idx3 = 0  # Current index in the results array

    # There is at least one element in each subarray:
    v1 = next(arr1)
    v2 = next(arr2)
    while idx3 < n_elements:
        if v1 < v2:
            results[idx3] = v1
            n1 -= 1 # One fewer element to process
            v1 = next(arr1) if n1 else inf
        else:
            results[idx3] = v2
            n2 -= 1
            v2 = next(arr2) if n2 else inf

        idx3 += 1

    return results

if __name__ == '__main__':
    print(solution([-3, -2, -1]))  # All negative
    print(solution([2, 4, 6]))  # All non-negative
    print(solution([-4, -2, 0, 3, 6, 9]))  # A mixture

Prints:

[1, 4, 9]
[4, 16, 36]
[0, 4, 9, 16, 36, 81]

Answer 3

It seems that you are sorting the elements of the array based on the absolute values of the elements. This actually creates loss of generality; forcing you to create cases, will certainly slow you down.

The more efficient method would be to square the elements inside the array first then apply a sorting algorithm, as this allows you to create comparison between the elements more easily, while fulfilling one of the conditions the question asks for; multitasking if you will.

As for your sorting algorithm, it may be best to use counting sort, working in linear time, specifically O(n + k). The counting sort only works for arrays with elements bound from [0, k].

Improved Method

• Loop through the array
• For each element, nums[i], let nums[i] = nums[i]**2 (square each element)
• Use counting sort (or any other linear time sorting algorithm)
• Return array

All steps involved run in linear time.

Answer 4

Portability

The versions of Python 3.x I tried were not happy with the function type hint List. This works for me:

def sortedSquares(self, nums):

Simpler

This line:

if nums[0] == 0 or nums[0] > 0:

can be simplified using a single comparison:

if nums[0] >= 0:

This may lead to a tiny increase in performance.

The same is true for:

elif nums[n-1] == 0 or nums[n-1] < 0:

Lines like this:

k = k+1

are simpler using the special assignment operators:

k += 1

Naming

Lower-case l is a bad name for a variable because it is easily confused with the number 1, not to mention that it conveys little meaning:

for l in range(s,q-1):

Layout

The black program can be used to automatically to add consistency to how you use whitespace, particularly around operators.

The code would be a little easier to read and understand with some blank lines between main sections in the function, such as around the while loops and if/else statements.

DRY

These line are repeated 3 times:

for i in range(n):
    nums[i] = nums[i]**2

It would be worth creating a helper function and calling that. I don't think it will hurt performance.

Comments

The comments are somewhat helpful in that they note the calculation complexity. It would also be helpful to add other comments on what the code is doing.

---

While these suggestions are unlikely to directly improve performance, they may lead to code that is easier to understand, perhaps making it easier to make the substantial modifications you need to reach the goal.

User @J_H identified this as LeetCode problem 977. I searched Stack Overflow, and found this related question on the python tag:

Squares of a Sorted Array: why sorted() method is faster than O(n) method?

Of particular interest is this Answer.

Answer 5

This code looks a lot like a C programmer's attempt to write Python. It's full of hand-built loops indexing the list, and doesn't make any use of Python's rich standard library.

The two functions that could replace most of the code we have are bisect.bisect() and heapq.merge(); another Python feature that helps is making slices of parts of the input list.

Here's how I'd use those functions to write an alternative:

from bisect import bisect
from heapq import merge

def square(n):
    return n*n

def sorted_squares(nums):
    z = bisect(nums, 0)  # zero-crossing point
    return merge(map(square, reversed(nums[:z])),
                 map(square, nums[z:]))

See how much easier this is to reason about, without having to hold the intricacies of control flow in your head as you read? Just three lines of list manipulation replaces 45 lines in the original, and we've reduced the number of variables from at least 6 down to just one. And instead of four levels of indentation, everything is now at a single level. All that goes together to make something that's much more comprehensible.

If memory use is vitally important, we might use itertools.islice() to avoid copying list contents (at the expense of reducing clarity somewhat):

    neg = islice(reversed(nums), len(nums)-z, None)
    pos = islice(nums, z, None)
    return merge(map(square, neg), map(square, pos))

---

Quick demo:

if __name__ == '__main__':
    print(list(sorted_squares([-4,-1,0,3,10])))

That gives the expected result:

[0, 1, 9, 16, 100]

---

And converting to list for the stated requirement:

class Solution:
    def sortedSquares(self, nums: list[int]) -> list[int]:
        return list(sorted_squares(nums))

I believe it's worth separating the conversion like that, if the function might be reused somewhere that constructing the whole list in memory is unnecessary.