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.

1. The sum of the respective array elements of A and B is equal.
2. 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

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]))