360 Sort Transformed Array

Given a sorted array of integers nums and integer values a,b and c. Apply a quadratic function of the form f(x) =ax2+bx+cto each element x in the array.

The returned array must be in sorted order.

Expected time complexity:O(n)

Example:

nums = [-4, -2, 2, 4], a = 1, b = 3, c = 5,

Result: [3, 9, 15, 33]

nums = [-4, -2, 2, 4], a = -1, b = 3, c = 5

Result: [-23, -5, 1, 7]

Complexity: O(n) time and constant extra space

import operator


class Solution:
    def sortTransformedArray(self, nums, a, b, c):
        """
        :type nums: List[int]
        :type a: int
        :type b: int
        :type c: int
        :rtype: List[int]
        """

        def quad(x):
            return (a * pow(x, 2)) + (b * x) + c

        solution = [0] * len(nums)
        def iter_forward_or_rev(start, add, cmp):
            left = 0
            right = len(nums) - 1
            while left <= right:
                if cmp(quad(nums[left]), quad(nums[right])):
                    solution[start] = quad(nums[left])
                    left += 1
                else:
                    solution[start] = quad(nums[right])
                    right -= 1
                start += add

        # if a > 0, then left of parabola are largest
        # otherwise the ends of parabola are the smallest
        if a > 0:
            iter_forward_or_rev(len(nums) - 1, -1, operator.gt)
        else:
            iter_forward_or_rev(0, 1, operator.lt)
        return solution

Last updated