BinarySearch 0801 Quadratic Application

Problem statement

https://binarysearch.com/problems/Quadratic-Application/

Solution

Equal to Leetcode 0360. Sort Transformed Array.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, nums, a, b, c):
        if a == 0 and b >= 0: return [b*x+c for x in nums]
        if a == 0 and b <  0: return [b*x+c for x in nums[::-1]]
        p1 = [a*x*x+b*x+c for x in nums if x < -b/2/a]
        p2 = [a*x*x+b*x+c for x in nums if x >= -b/2/a]
        if a > 0: return sorted(p1[::-1] + p2)
        if a < 0: return sorted(p2[::-1] + p1)