2948. Make Lexicographically Smallest Array by Swapping Elements

• Time: $O(n\log n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  vector<int> lexicographicallySmallestArray(vector<int>& nums, int limit) {
    vector<int> ans(nums.size());
    // [[(num, index)]], where the difference between in each pair in each
    // `[(num, index)]` group <= `limit`
    vector<vector<pair<int, int>>> numAndIndexesGroups;

    for (const pair<int, int>& numAndIndex : getNumAndIndexes(nums))
      if (numAndIndexesGroups.empty() ||
          numAndIndex.first - numAndIndexesGroups.back().back().first > limit) {
        // Start a new group.
        numAndIndexesGroups.push_back({numAndIndex});
      } else {
        // Append to the existing group.
        numAndIndexesGroups.back().push_back(numAndIndex);
      }

    for (const vector<pair<int, int>>& numAndIndexesGroup :
         numAndIndexesGroups) {
      vector<int> sortedNums;
      vector<int> sortedIndices;
      for (const auto& [num, index] : numAndIndexesGroup) {
        sortedNums.push_back(num);
        sortedIndices.push_back(index);
      }
      ranges::sort(sortedIndices);
      for (int i = 0; i < sortedNums.size(); ++i)
        ans[sortedIndices[i]] = sortedNums[i];
    }

    return ans;
  }

 private:
  vector<pair<int, int>> getNumAndIndexes(const vector<int>& nums) {
    vector<pair<int, int>> numAndIndexes;
    for (int i = 0; i < nums.size(); ++i)
      numAndIndexes.emplace_back(nums[i], i);
    ranges::sort(numAndIndexes);
    return numAndIndexes;
  }
};

class Solution {
  public int[] lexicographicallySmallestArray(int[] nums, int limit) {
    int[] ans = new int[nums.length];
    List<List<Pair<Integer, Integer>>> numAndIndexesGroups = new ArrayList<>();

    for (Pair<Integer, Integer> numAndIndex : getNumAndIndexes(nums))
      if (numAndIndexesGroups.isEmpty() ||
          numAndIndex.getKey() -
                  numAndIndexesGroups.get(numAndIndexesGroups.size() - 1)
                      .get(numAndIndexesGroups.get(numAndIndexesGroups.size() - 1).size() - 1)
                      .getKey() >
              limit) {
        // Start a new group.
        numAndIndexesGroups.add(new ArrayList<>(List.of(numAndIndex)));
      } else {
        // Append to the existing group.
        numAndIndexesGroups.get(numAndIndexesGroups.size() - 1).add(numAndIndex);
      }

    for (List<Pair<Integer, Integer>> numAndIndexesGroup : numAndIndexesGroups) {
      List<Integer> sortedNums = new ArrayList<>();
      List<Integer> sortedIndices = new ArrayList<>();
      for (Pair<Integer, Integer> pair : numAndIndexesGroup) {
        sortedNums.add(pair.getKey());
        sortedIndices.add(pair.getValue());
      }
      sortedIndices.sort(null);
      for (int i = 0; i < sortedNums.size(); ++i) {
        ans[sortedIndices.get(i)] = sortedNums.get(i);
      }
    }

    return ans;
  }

  private Pair<Integer, Integer>[] getNumAndIndexes(int[] nums) {
    Pair<Integer, Integer>[] numAndIndexes = new Pair[nums.length];
    for (int i = 0; i < nums.length; ++i)
      numAndIndexes[i] = new Pair<>(nums[i], i);
    Arrays.sort(numAndIndexes, Comparator.comparing(Pair::getKey));
    return numAndIndexes;
  }
}

class Solution:
  def lexicographicallySmallestArray(
      self,
      nums: list[int],
      limit: int,
  ) -> list[int]:
    ans = [0] * len(nums)
    numAndIndexes = sorted([(num, i) for i, num in enumerate(nums)])
    # [[(num, index)]], where the difference between in each pair in each
    # `[(num, index)]` group <= `limit`
    numAndIndexesGroups: list[list[tuple[int, int]]] = []

    for numAndIndex in numAndIndexes:
      if (not numAndIndexesGroups or
              numAndIndex[0] - numAndIndexesGroups[-1][-1][0] > limit):
        # Start a new group.
        numAndIndexesGroups.append([numAndIndex])
      else:
        # Append to the existing group.
        numAndIndexesGroups[-1].append(numAndIndex)

    for numAndIndexesGroup in numAndIndexesGroups:
      sortedNums = [num for num, _ in numAndIndexesGroup]
      sortedIndices = sorted([index for _, index in numAndIndexesGroup])
      for num, index in zip(sortedNums, sortedIndices):
        ans[index] = num

    return ans