LeetCode 805. Split Array With Same Average

September 02, 2020

the lists don't need to be evenly split
we could find the average of the full input and then return true iff there exists a strict subset that has the same average
it sounds like this reduces to a subset sum problem

This solution was implemented with modified versions of wikipedia's pseudo-polynomial and exponential subset sum algorithms. Both algorithms now also store the number of elements in each subset in order to:

make sure we don't create a subset that's the same size as $A$ in the exponential solution
compute averages in the dynamic programming solution

We decide which algorithm to use by comparing expected run times based on the sizes of $n$ and $max(A)$ .

from fractions import Fraction
from itertools import chain, combinations
import math
import random
from typing import List
import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_examples(self):
        self.assertTrue(self.s.splitArraySameAverage([1,2,3,4,5,6,7,8]))
        self.assertFalse(self.s.splitArraySameAverage([3,1]))

    def test_empty(self):
        self.assertFalse(self.s.splitArraySameAverage([0]))

    def test_all_zeros(self):
        for i in range(2,6):
            self.assertTrue(self.s.splitArraySameAverage([0]*i))

    def test_unsorted(self):
        a = [1,2,3,4,5,6,7,8]
        for i in range(10):
            random.shuffle(a)
            self.assertTrue(self.s.splitArraySameAverage(a))

    def test_small_false(self):
        self.assertFalse(self.s.splitArraySameAverage([0,12,3,1]))
        self.assertFalse(self.s.splitArraySameAverage([0,1,2,3,1]))

    def test_large_values(self):
        self.assertFalse(self.s.splitArraySameAverage([3863,703,1799,327,3682,4330,3388,6187,5330,6572,938,6842,678,9837,8256,6886,2204,5262,6643,829,745,8755,3549,6627,1633,4290,7]))
        self.assertFalse(self.s.splitArraySameAverage([1408,3639,3116,4391,1150,827,8914,8567,9471,3438,8731,5405,7878,5702,9]))

    def test_large_n(self):
        self.assertTrue(self.s.splitArraySameAverage([0,1,2,3]*100))

class Solution:

    def splitArraySameAverage(self, A: List[int]) -> bool:
        n = len(A)
        s = max(A)
        if s == 0 or math.log(s*n, 2) < n/2:
            return self.pseudo_polynomial(A)
        return self.exponential(A)

    def exponential(self, A: List[int]) -> bool:
        def powerset(iterable):
            # from: https://docs.python.org/3/library/itertools.html#itertools-recipes
            s = list(iterable)
            return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
        if len(A) <= 1:
            return False
        x = sorted(A)
        x = [v - Fraction(sum(x), len(x)) for v in x]
        a = sorted([(sum(s), len(s)) for s in powerset(x[:math.floor(len(x)/2)]) if len(s) > 0])
        if 0 in a:
            return True
        b = sorted([(sum(s), len(s)) for s in powerset(x[math.floor(len(x)/2):]) if len(s) > 0])
        if 0 in b:
            return True
        i = len(a)-1
        j = 0
        while i >= 0 and j < len(b):
            cs = a[i][0] + b[j][0]
            if cs == 0:
                if a[i][1] + b[j][1] < len(x):
                    return True
                if j == len(b)-1:
                    i -= 1
                elif i == 0:
                    j += 1
            elif cs > 0:
                i -= 1
            else:
                j += 1
        return False

    def pseudo_polynomial(self, A: List[int]) -> bool:
        if len(A) <= 1:
            return False
        x = sorted(A)
        x_avg = sum(x) / len(x)
        q = {}
        for i in range(1, len(x)+1):
            b = math.ceil(x_avg * i)
            self.Q(x[:i], 0, b, q)
        for (i,s) in q:
            if q[(i,s)][0] and q[(i,s)][1] < len(x) and s / q[(i,s)][1] == x_avg:
                # If:
                # * there is a subset in x[:i+1] that sums to s
                # * the size of that subset is smaller than the size of x
                # * the average of that subset is the same as the average of x
                return True
        return False

    def Q(self, x: List[int], a: int, b: int, q: dict) -> int:
        i = len(x)
        for s in range(a, b+1):
            q[(i,s)] = (x[i-1] == s, 1)
            if i == 1:
                continue
            q_p_1 = (i-1,s)
            if q_p_1 in q and q[q_p_1][0]:
                q[(i,s)] = q[q_p_1]
            q_p_2 = (i-1,s-x[i-1])
            if q_p_2 in q and q[q_p_2][0]:
                q[(i,s)] = (True, q[q_p_2][1] + 1)

if __name__ == "__main__":
    unittest.main()