Program to find squared elements list in sorted order in Python

When we have a sorted array with negative and positive numbers, squaring them can disrupt the order. We need to find squared elements and return them in sorted order efficiently.

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

Algorithm

To solve this efficiently in O(n) time, we use a two-pointer approach ?

• Initialize left pointer at start (0) and right pointer at end (n-1)
• Create result array and fill from the end (largest squared values first)
• Compare absolute values at both pointers
• Square the larger absolute value and place it at current position
• Move the corresponding pointer and decrease index
• Continue until all elements are processed

Example

Let us see the implementation to get better understanding ?

def solve(nums):
    n = len(nums)
    left = 0
    right = n - 1
    index = n - 1
    result = [0] * n
    
    while index >= 0:
        if abs(nums[left]) > abs(nums[right]):
            result[index] = nums[left] * nums[left]
            left += 1
        else:
            result[index] = nums[right] * nums[right]
            right -= 1
        index -= 1
    
    return result

# Test the function
nums = [-8, -3, 0, 5, 6]
print("Input:", nums)
print("Output:", solve(nums))

Input: [-8, -3, 0, 5, 6]
Output: [0, 9, 25, 36, 64]

How It Works

The algorithm works by recognizing that the largest squared values come from either the most negative or most positive numbers. By comparing absolute values from both ends, we can place squares in descending order from the end of the result array ?

# Step-by-step trace
nums = [-8, -3, 0, 5, 6]

# Step 1: |?8| vs |6| ? 8 > 6, so 64 goes to result[4]
# Step 2: |?3| vs |6| ? 3 < 6, so 36 goes to result[3]  
# Step 3: |?3| vs |5| ? 3 < 5, so 25 goes to result[2]
# Step 4: |?3| vs |0| ? 3 > 0, so 9 goes to result[1]
# Step 5: |0| only left, so 0 goes to result[0]

def solve_with_trace(nums):
    n = len(nums)
    left, right, index = 0, n - 1, n - 1
    result = [0] * n
    
    while index >= 0:
        left_abs = abs(nums[left])
        right_abs = abs(nums[right])
        
        if left_abs > right_abs:
            result[index] = nums[left] ** 2
            print(f"Placing {nums[left]}² = {result[index]} at index {index}")
            left += 1
        else:
            result[index] = nums[right] ** 2
            print(f"Placing {nums[right]}² = {result[index]} at index {index}")
            right -= 1
        index -= 1
    
    return result

nums = [-4, -1, 0, 3, 10]
print("Tracing steps:")
result = solve_with_trace(nums)
print("Final result:", result)

Tracing steps:
Placing 10² = 100 at index 4
Placing (-4)² = 16 at index 3
Placing 3² = 9 at index 2
Placing (-1)² = 1 at index 1
Placing 0² = 0 at index 0
Final result: [0, 1, 9, 16, 100]

Time and Space Complexity

Conclusion

The two-pointer approach efficiently squares and sorts elements in O(n) time by comparing absolute values from both ends. This avoids the need for sorting after squaring, making it optimal for sorted input arrays.