Problem 2966

Question statement:

You are given an integer array nums of size n and a positive integer k. Divide the array into one or more arrays of size 3 satisfying the conditions below, and return a 2D array containing all the arrays. If it is impossible to satisfy the conditions, return an empty array. And if there are multiple answers, return any of them.

• Each element of nums should be in exactly one array.
• The difference between any two elements in one array is less than or equal to k.

Solution

Main idea:

Sort the array first. Scan through the array and batch elements into groups of three. For any given batch, if at any point batch[j] - batch[0] > k, then we simply return [] since this is a violation of item 2 (and all the further elements in future batches are by design bigger than batch[j]). Otherwise append batch[j] into a temp array, and then once the batch is complete, add the list to the ans array, and move onto the next batch.

Code:

def divideArray(self, nums: List[int], k: int) -> List[List[int]]:
    nums = sorted(nums)
    ans = []
    pointer = 0
    while pointer < len(nums):
        temp = [nums[pointer]]
        for j in range(1,3):
            if nums[pointer+j] - temp[0] <= k:
                temp.append(nums[pointer+j])
            else:
                return []
        pointer += 3
        ans.append(temp)
    return ans

Comments:

Time complexity: O(nlog n) for sorting.

At the time of submission, the above solution passes leetcode submission cases in 741ms, beats 85.65% of users. It uses 31.88MB of memory, beating 29.03% of users.