leetcode - 2948. Make Lexicographically Smallest Array by Swapping Elements

Description

You are given a 0-indexed array of positive integers nums and a positive integer limit.

In one operation, you can choose any two indices i and j and swap nums[i] and nums[j] if |nums[i] - nums[j]| <= limit.

Return the lexicographically smallest array that can be obtained by performing the operation any number of times.

An array a is lexicographically smaller than an array b if in the first position where a and b differ, array a has an element that is less than the corresponding element in b. For example, the array [2,10,3] is lexicographically smaller than the array [10,2,3] because they differ at index 0 and 2 < 10.

Example 1:

Input: nums = [1,5,3,9,8], limit = 2
Output: [1,3,5,8,9]
Explanation: Apply the operation 2 times:
- Swap nums[1] with nums[2]. The array becomes [1,3,5,9,8]
- Swap nums[3] with nums[4]. The array becomes [1,3,5,8,9]
We cannot obtain a lexicographically smaller array by applying any more operations.
Note that it may be possible to get the same result by doing different operations.

Example 2:

Input: nums = [1,7,6,18,2,1], limit = 3
Output: [1,6,7,18,1,2]
Explanation: Apply the operation 3 times:
- Swap nums[1] with nums[2]. The array becomes [1,6,7,18,2,1]
- Swap nums[0] with nums[4]. The array becomes [2,6,7,18,1,1]
- Swap nums[0] with nums[5]. The array becomes [1,6,7,18,1,2]
We cannot obtain a lexicographically smaller array by applying any more operations.

Example 3:

Input: nums = [1,7,28,19,10], limit = 3
Output: [1,7,28,19,10]
Explanation: [1,7,28,19,10] is the lexicographically smallest array we can obtain because we cannot apply the operation on any two indices.

Constraints:

1 <= nums.length <= 10^5
1 <= nums[i] <= 10^9
1 <= limit <= 10^9

Solution

If a could swap with b, b could swap with c, then a could swap with c. So use a sliding window to union all sets, and then start with the smallest number, find its final location.

Time complexity: $o(n)$
Space complexity: $o(n)$

Code

class UnionFindSet:
    def __init__(self, nums: list) -> None:
        self.parent = [i for i in nums]
        self.height = [1] * len(nums)

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int) -> None:
        x, y = self.find(x), self.find(y)
        if x == y:
            return 
        if self.height[x] < self.height[y]:
            self.parent[x] = y
        else:
            self.parent[y] = x
            if self.height[x] == self.height[y]:
                self.height[x] += 1

    def same(self, x: int, y: int) -> bool:
        return self.find(x) == self.find(y)


class Solution:
    def lexicographicallySmallestArray(self, nums: List[int], limit: int) -> List[int]:
        sorted_num = sorted([(index, item) for index, item in enumerate(nums)], key=lambda x: x[1])
        sorted_num_val = sorted(nums)
        union_find = UnionFindSet(list(range(len(nums))))
        # sliding window to union all sets
        left = 0
        for right in range(len(sorted_num)):
            while sorted_num[right][1] - sorted_num[left][1] > limit:
                left += 1
            union_find.union(sorted_num[right][0], sorted_num[left][0])
        swap_map = {}
        # key: parent index, value: children indexs (indexs that could be swapped)
        for i in range(len(nums)):
            parent_i = union_find.find(i)
            if parent_i not in swap_map:
                swap_map[parent_i] = []
            heapq.heappush(swap_map[parent_i], i)
        # start from the smallest number
        res = [0] * len(nums)
        for i in range(len(sorted_num)):
            index_in_num = sorted_num[i][0]
            candidates = swap_map[union_find.find(index_in_num)]
            target_index = heapq.heappop(candidates)
            res[target_index] = sorted_num[i][1]
        return res