Determine whether any permutation of a given array exists such that the sum of all subarrays of length K are equal

Question

I was asked this question in an interview.

You are given an array consisting of N integers. You need to determine whether any permutation of a given array exists such that the sum of all subarrays of length K are equal. N is divisible by K i.e. N mod K = 0.
If the array is [1,2,1,3,2,3] and K = 2, it is False.
If the array is [1,2,1,3,2,3] and K = 3, it is True.

My logic was this.

1. If N == K, return True.
2. Else store each of the distinct element in the array in a dictionary. If the number of distinct elements in the dictionary is greater than K, return False.
3. Else store the count of each distinct element in the dictionary.
4. If the count of any distinct element is less than (N/K), return False.
5. Finally return True.

Here is my code in Python:

def check(nums, n, k):
if n == k:
    return True
else:
    my_dict = {}
    for i in nums:
        if i not in my_dict:
            my_dict[i] = 1
        else:
            my_dict[i] += 1
    if len(my_dict) > k:
        return False
    count = int(n/k)
    for i,j in my_dict.items():
        if j < count:
            return False
    return True  

Am I doing the right approach ?

Answer 1

Use the collections.Counter instead of keeping a count yourself.

When you have counted all the elements from the list, check if any of their counts is not a factor of \$ N \$ and return False.

Therefore:

from collections import Counter

def check(nums, n, k):
    if n == k:
        return True
    counts = Counter(nums)
    if len(counts) > k:
        return False
    return all((n % x == 0 for x in counts.values()))