Taro Logo

Partition Array According to Given Pivot

Medium
Amazon logo
Amazon
1 view
Topics:
Arrays

You are given a 0-indexed integer array nums and an integer pivot. Rearrange nums such that the following conditions are satisfied:

  • Every element less than pivot appears before every element greater than pivot.
  • Every element equal to pivot appears in between the elements less than and greater than pivot.
  • The relative order of the elements less than pivot and the elements greater than pivot is maintained.
    • More formally, consider every pi, pj where pi is the new position of the ith element and pj is the new position of the jth element. If i < j and both elements are smaller (or larger) than pivot, then pi < pj.

Return nums after the rearrangement.

Example 1:

Input: nums = [9,12,5,10,14,3,10], pivot = 10
Output: [9,5,3,10,10,12,14]
Explanation: 
The elements 9, 5, and 3 are less than the pivot so they are on the left side of the array.
The elements 12 and 14 are greater than the pivot so they are on the right side of the array.
The relative ordering of the elements less than and greater than pivot is also maintained. [9, 5, 3] and [12, 14] are the respective orderings.

Example 2:

Input: nums = [-3,4,3,2], pivot = 2
Output: [-3,2,4,3]
Explanation: 
The element -3 is less than the pivot so it is on the left side of the array.
The elements 4 and 3 are greater than the pivot so they are on the right side of the array.
The relative ordering of the elements less than and greater than pivot is also maintained. [-3] and [4, 3] are the respective orderings.

Constraints:

  • 1 <= nums.length <= 105
  • -106 <= nums[i] <= 106
  • pivot equals to an element of nums.

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. What is the data type of the elements in the array and the pivot value? Can I assume they are integers?
  2. Can the input array be empty or null? If so, what should I return in those cases?
  3. Are there any constraints on the range of values within the array?
  4. If there are elements equal to the pivot, where should they be placed in the partitioned array - immediately before or after the pivot-smaller elements?
  5. Can I modify the input array in-place, or should I return a new array with the partitioned elements?

Brute Force Solution

Approach

The brute force approach to partitioning an array around a pivot involves checking every possible arrangement of the numbers. We will separate the numbers into three groups based on their comparison to the pivot value. We'll try every single possible ordering to find the one that satisfies the partition requirement.

Here's how the algorithm would work step-by-step:

  1. Imagine we have three empty containers labeled 'less than pivot', 'equal to pivot', and 'greater than pivot'.
  2. Take the first number from the original list.
  3. Try putting that number in the first container ('less than pivot'). Then, try putting it in the second container ('equal to pivot'). Finally, try putting it in the third container ('greater than pivot').
  4. For each of those placements, take the next number from the original list and again try placing it in each of the three containers.
  5. Continue this process until all numbers from the original list have been placed into the containers.
  6. Each time we've placed all the numbers, we check if the containers are arranged in the correct order: 'less than pivot' first, followed by 'equal to pivot', and finally 'greater than pivot'.
  7. If the order is correct, we keep track of this arrangement.
  8. After trying all possible arrangements, we pick one of the correct arrangements as our solution. If we are asked to keep the original order of elements within each partition, the brute force approach would need to enumerate all permutations while maintaining the original relative order within each group, which would be highly inefficient.

Code Implementation

def partition_array_brute_force(numbers, pivot):
    all_arrangements = []

    def generate_arrangements(index, less_than, equal_to, greater_than):
        if index == len(numbers):
            # We've placed all numbers, check order
            all_arrangements.append((less_than, equal_to, greater_than))
            return

        current_number = numbers[index]

        # Try putting the number in 'less than'
        generate_arrangements(index + 1, less_than + [current_number], equal_to, greater_than)

        # Try putting the number in 'equal to'
        generate_arrangements(index + 1, less_than, equal_to + [current_number], greater_than)

        # Try putting the number in 'greater than'
        generate_arrangements(index + 1, less_than, equal_to, greater_than + [current_number])

    generate_arrangements(0, [], [], [])

    for less_than_numbers, equal_to_numbers, greater_than_numbers in all_arrangements:
        
        #Check if the current arrangement is correctly partitioned
        is_correctly_partitioned = True

        reconstructed_array = less_than_numbers + equal_to_numbers + greater_than_numbers

        less_than_group = [num for num in reconstructed_array if num < pivot]
        equal_to_group = [num for num in reconstructed_array if num == pivot]
        greater_than_group = [num for num in reconstructed_array if num > pivot]

        reconstructed_array_check = less_than_group + equal_to_group + greater_than_group
        if reconstructed_array != reconstructed_array_check:
            is_correctly_partitioned = False
        
        if is_correctly_partitioned:
            #Ensures original order is considered when reconstructing the array
            result = less_than_numbers + equal_to_numbers + greater_than_numbers
            return result

    return []

Big(O) Analysis

Time Complexity
O(3^n * n)The described brute force approach explores all possible placements of each element into one of three containers (less than, equal to, greater than the pivot). For each of the n elements, we have 3 choices, leading to 3^n possible arrangements. After generating each arrangement, we need to verify if it satisfies the partitioning condition, which takes O(n) time to iterate through the arrangement and validate its correctness. Therefore, the overall time complexity is O(3^n * n).
Space Complexity
O(N)The brute force approach, as described, uses three containers ('less than pivot', 'equal to pivot', and 'greater than pivot') to store the partitioned elements. In the worst-case scenario, all N elements of the input array could end up in one of these containers. Therefore, the auxiliary space required to store these containers grows linearly with the input size N, resulting in O(N) space complexity.

Optimal Solution

Approach

The best way to rearrange the numbers around a specific value is to create three separate groups: numbers smaller than the value, numbers equal to the value, and numbers bigger than the value. Then, we simply put these groups back together in the correct order to get the final result.

Here's how the algorithm would work step-by-step:

  1. First, go through all the numbers and make three new temporary piles based on whether they are less than, equal to, or greater than the given value.
  2. After sorting all numbers into these piles, put the piles back together in order: the 'less than' pile first, followed by the 'equal to' pile, and finally the 'greater than' pile.
  3. This combined pile is the rearranged sequence that fits the requirement.

Code Implementation

def partition_array(numbers, pivot_value):
    less_than_pivot = []
    equal_to_pivot = []
    greater_than_pivot = []

    # Segregate numbers into three lists
    for number in numbers:
        if number < pivot_value:
            less_than_pivot.append(number)
        elif number == pivot_value:
            equal_to_pivot.append(number)
        else:
            greater_than_pivot.append(number)

    # Combining the three lists into a single result
    return less_than_pivot + equal_to_pivot + greater_than_pivot

def partition_array_inplace(numbers, pivot_value):
    less_than = []
    equal_to = []
    greater_than = []

    # Create 3 lists based on element's comparison to pivot.
    for number in numbers:
        if number < pivot_value:
            less_than.append(number)
        elif number == pivot_value:
            equal_to.append(number)
        else:
            greater_than.append(number)

    # Reconstruct the original array in place.
    numbers[:] = less_than + equal_to + greater_than

def partition_array_optimal(numbers, pivot_value):
    smaller_numbers = []
    equal_numbers = []
    larger_numbers = []

    # Separate numbers to smaller, equal, and larger groups.
    for number in numbers:
        if number < pivot_value:
            smaller_numbers.append(number)
        elif number == pivot_value:
            equal_numbers.append(number)
        else:
            larger_numbers.append(number)

    # First create smaller list,
    partitioned_array = smaller_numbers

    # Then the equal list,
    partitioned_array.extend(equal_numbers)

    # Finally add the larger list.
    partitioned_array.extend(larger_numbers)

    return partitioned_array

Big(O) Analysis

Time Complexity
O(n)The provided solution involves iterating through the input array of size n once to categorize elements into three separate groups (less than, equal to, and greater than the pivot). Subsequently, these groups are concatenated together. The dominant operation is the initial iteration through the array, which is directly proportional to the number of elements n. Therefore, the time complexity is O(n).
Space Complexity
O(N)The provided solution uses three temporary lists to store elements less than, equal to, and greater than the pivot. In the worst case, each of these lists could potentially hold all N elements of the input array if all elements fall into the same category (e.g., all elements are less than the pivot), or they are evenly distributed across the three lists. Therefore, the auxiliary space required grows linearly with the input size N. This results in a space complexity of O(N).

Edge Cases

CaseHow to Handle
Null or empty input arrayReturn an empty array or throw an exception, depending on problem specifications.
Array with a single elementReturn the array as is, since it is already partitioned.
Pivot is smaller than all elements in the arrayThe array will be partitioned into elements smaller than pivot followed by elements equal/greater than pivot.
Pivot is larger than all elements in the arrayThe array will be partitioned into elements smaller than pivot followed by elements equal/greater than pivot.
Array contains many duplicate elements equal to the pivotEnsure the algorithm places all pivot values in the 'equal' partition correctly.
Array contains negative numbers, zeros, and positive numbersThe partitioning logic should handle all number types correctly based on their comparison with the pivot.
Large array size exceeding available memoryConsider an in-place partitioning algorithm to minimize memory usage if applicable.
Pivot is a floating-point number and the array contains integers.Ensure proper type conversion and comparison between the pivot and array elements.