Find Squared Elements in Sorted Order in Python

Suppose we have a list of numbers called nums, where elements are sorted in ascending order, we have to square the elements and return the result in sorted order.

So, if the input is like nums = [-8, -3, 0, 5, 6], then the output will be [0, 9, 25, 36, 64]

To solve this, we will follow these steps −

• n := size of nums
• l := 0
• r := n - 1
• index := n - 1
• res := a list of size same as nums and fill it with 0
• while index >= 0, do
  • if |nums[l]| > |nums[r]|, then
    • res[index] := nums[l] * nums[l]
    • l := l + 1
  • otherwise,
    • res[index] := nums[r] * nums[r]
    • r := r - 1
  • index := index - 1
• return res

Example

Let us see the following implementation to get better understanding −

def solve(nums):
   n = len(nums)
   l = 0
   r = n - 1
   index = n - 1
   res = [0 for i in range(len(nums))]
   while index >= 0:
      if abs(nums[l]) > abs(nums[r]):
         res[index] = nums[l] * nums[l]
         l += 1
      else:
         res[index] = nums[r] * nums[r]
         r -= 1
      index -= 1

   return res

nums = [-8, -3, 0, 5, 6]
print(solve(nums))

Input

[-8, -3, 0, 5, 6]

Output

[0, 9, 25, 36, 64]