Skip to content

2161. Partition Array According to Given Pivot

SourceBiweekly Contest 71 Q2DifficultyMediumRating1337

Description

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.

Solutions

Solution 1: Simulation

Thinking

Partition relative to \(\textit{pivot}\) into less, equal, and greater parts, preserving order inside each part. A stable one-pass split suffices.

Collect three lists in encounter order and concatenate them.

Extra memory is linear.

We can traverse the array \(\textit{nums}\), sequentially finding all elements less than \(\textit{pivot}\), all elements equal to \(\textit{pivot}\), and all elements greater than \(\textit{pivot}\), then concatenate them in the order required by the problem.

Time complexity \(O(n)\), where \(n\) is the length of the array \(\textit{nums}\). Ignoring the space consumption of the answer array, the space complexity is \(O(1)\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        a, b, c = [], [], []
        for x in nums:
            if x < pivot:
                a.append(x)
            elif x == pivot:
                b.append(x)
            else:
                c.append(x)
        return a + b + c

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution {
    public int[] pivotArray(int[] nums, int pivot) {
        int n = nums.length;
        int[] ans = new int[n];
        int k = 0;
        for (int x : nums) {
            if (x < pivot) {
                ans[k++] = x;
            }
        }
        for (int x : nums) {
            if (x == pivot) {
                ans[k++] = x;
            }
        }
        for (int x : nums) {
            if (x > pivot) {
                ans[k++] = x;
            }
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
public:
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        vector<int> ans;
        for (int& x : nums) {
            if (x < pivot) {
                ans.push_back(x);
            }
        }
        for (int& x : nums) {
            if (x == pivot) {
                ans.push_back(x);
            }
        }
        for (int& x : nums) {
            if (x > pivot) {
                ans.push_back(x);
            }
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
func pivotArray(nums []int, pivot int) []int {
    var ans []int
    for _, x := range nums {
        if x < pivot {
            ans = append(ans, x)
        }
    }
    for _, x := range nums {
        if x == pivot {
            ans = append(ans, x)
        }
    }
    for _, x := range nums {
        if x > pivot {
            ans = append(ans, x)
        }
    }
    return ans
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
function pivotArray(nums: number[], pivot: number): number[] {
    const ans: number[] = [];
    for (const x of nums) {
        if (x < pivot) {
            ans.push(x);
        }
    }
    for (const x of nums) {
        if (x === pivot) {
            ans.push(x);
        }
    }
    for (const x of nums) {
        if (x > pivot) {
            ans.push(x);
        }
    }
    return ans;
}

Solution 2: Two pointers

Thinking

Solution 1 uses three buffers. Initializing the answer with \(\textit{pivot}\) avoids writing the equal part.

Fill lesser values from the left and greater values from the right; the middle stays equal. Two opposing scans keep each side’s relative order.

This fill-in-place variant is shown in TypeScript / JavaScript.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
function pivotArray(nums: number[], pivot: number): number[] {
    const n = nums.length;
    const res = Array(n).fill(pivot);

    for (let i = 0, l = 0, j = n - 1, r = n - 1; i < n; i++, j--) {
        if (nums[i] < pivot) res[l++] = nums[i];
        if (nums[j] > pivot) res[r--] = nums[j];
    }

    return res;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
function pivotArray(nums, pivot) {
    const n = nums.length;
    const res = Array(n).fill(pivot);

    for (let i = 0, l = 0, j = n - 1, r = n - 1; i < n; i++, j--) {
        if (nums[i] < pivot) res[l++] = nums[i];
        if (nums[j] > pivot) res[r--] = nums[j];
    }

    return res;
}

Comments