A more efficient solution for balanced split of an array with additional conditions

Question

I appreciate your help in advance. This is a practice question from Meta’s interview preparation website. I have solved it, but I wonder if any optimization can be done.

Question:

Is there a way to solve the following problem with a time complexity of O(n)?

Problem Description:

You have been given an array nums of type int. Write a program that
returns the bool type as the return value of the solution function to
determine whether it is possible to split nums into two arrays A and B
such that the following two conditions are satisfied.

The sum of the respective array elements of A and B is equal.

All the array elements in A are strictly smaller than all the array elements in B.

Examples:

nums = [1,5,7,1] -> true since A = [1,1,5], B = [7]

nums = [12,7,6,7,6] -> false since A = [6,6,7], B = [7,12] failed the 2nd
requirement

What I have tried:

I have used the sort function in Python, which has a time complexity of O(nlog(n)).

from typing import List



def solution(nums: List[int]) -> bool:
    total_sum = sum(nums)

    # If the total sum of the array is 0 or an odd number, 
    # it is impossible to have array A and array B equal.
    if total_sum % 2 or total_sum == 0:
        return False

    nums.sort()

    curr_sum, i = total_sum, 0

    while curr_sum > total_sum // 2:
        curr_sum -= nums[i]
        i += 1

    if curr_sum == total_sum // 2 and nums[i] != nums[i - 1]:
        return True

    return False

Asked By: コリン

||

Source

Answer 1

For what it’s worth, you can modify QuickSelect to get a with-high-probability and expected O(n)-time algorithm, though Python’s sort is so fast that it hardly seems like a good idea. Deterministic O(n) is possible and left as an easy exercise to the reader familiar with selection algorithms (but the constant factor is even worse, so…).

import random


def can_partition(nums, a_sum=0, b_sum=0):
    if not nums:
        # True makes more sense here, but whatever
        return False
    pivot = random.choice(nums)
    less = sum(n for n in nums if n < pivot)
    equal = sum(n for n in nums if n == pivot)
    greater = sum(n for n in nums if n > pivot)
    a_ext = a_sum + less
    b_ext = greater + b_sum
    if abs(a_ext - b_ext) == equal:
        return True
    elif a_ext < b_ext:
        return can_partition([n for n in nums if n > pivot], a_ext + equal, b_sum)
    else:
        return can_partition([n for n in nums if n < pivot], a_sum, equal + b_ext)


print(can_partition([1, 5, 7, 1]))
print(can_partition([12, 7, 6, 7, 6]))

Answered By: David Eisenstat

Answer 2

Quickselect can potentially bring down the time complexity to O(n) just like David Eisenstat said… However, beyond algorithm consideration, sorting is probably faster than quickselect in real practice.

Other algorithms like min-heap and max-heap could also work, but the time complexity would be O(nlog(n)).

Answered By: namedtuple

A more efficient solution for balanced split of an array with additional conditions

Question:

Question:

Problem Description:

Examples:

What I have tried:

Answers: