Maximum sum of distinct subarrays with k stuck at flag logic

Question

Problem

You've given an integer array nums and integer k. Find the maximum subarray sum of all the subarraay of nums that meet the following conditions.

• The length of the subarray is k, and
• All the elements of the subarray are distinct.

If no subarray meets the conditions, return 0.

Link to the problem

Thought process

Keep a flag to know when its sure to compare. Update the flag to false when theres multiple frequencies using a dict. The issue is that if there are multiple discrepencies then i dont know when to turn back the flag to true? I am a little lost on that. Any idea would help me. Specifically i can't pass the third case in the tests.

Code

class Solution:
    def maximumSubarraySum(self, a, k):
        sumSoFar, toReturn, i, j = 0, 0, 0, 0
        canIcount = True
        eleToFreq = dict()
        N = len(a)

        if len(set(a)) < k:  # Handle all-duplicate case
            return 0

        while j < N:
            if j >= k - 1:
                if a[j] in eleToFreq:
                    eleToFreq[a[j]] += 1
                    canIcount = False
                else:
                    eleToFreq[a[j]] = 1
                sumSoFar += a[j]
                if canIcount:
                    toReturn = max(toReturn, sumSoFar)
                #print('does ' , eleToFreq, ' have ', a[i])
                if eleToFreq[a[i]] >= 1:
                    eleToFreq[a[i]] -= 1
                else:
                    eleToFreq.pop(a[i])
                sumSoFar-=a[i]
                i+=1
                j+=1
            else:
                if a[j] in eleToFreq:
                    eleToFreq[a[j]] += 1
                    canIcount = False
                else:
                    eleToFreq[a[j]] = 1
                sumSoFar += a[j]
                j+=1

        return toReturn


solution = Solution()
arr = [1, 5, 4, 2, 9, 9, 9]
k = 3
a = solution.maximumSubarraySum(arr, k)
print("Output:", a)
# Expected output: 15 actual output: 15

arr = [4, 4, 4]
k = 3
a = solution.maximumSubarraySum(arr, k)
print("Output:", a)
# Expected output: 0 actual output: 0

arr = [1,1,1,7,8,9]
k = 3
a = solution.maximumSubarraySum(arr, k)
print("Output:", a)
# Expected output: 24 actual output:0

Answer 1

The issue is that if there are multiple discrepancies then I don't know when to turn back the flag to true?

Your flag canIcount would need to be set to true again when all values in the current window are unique. This happens when the count of distinct values in the window is k, in other words when the number of keys in your dictionary is k. And as Kelly pointed out in comments this means you can do without that flag, and just check len(eleToFreq) == k.

Another issue is the condition eleToFreq[a[i]] >= 1, which should be eleToFreq[a[i]] > 1 as in the other case you'd want to pop the entry.

With those two things adapted with minimal changes to your code, we get this working code:

class Solution:
    def maximumSubarraySum(self, a, k):
        sumSoFar, toReturn, i, j = 0, 0, 0, 0
        eleToFreq = dict()
        N = len(a)

        if len(set(a)) < k:
            return 0

        while j < N:
            if j >= k - 1:
                if a[j] in eleToFreq:
                    eleToFreq[a[j]] += 1
                else:
                    eleToFreq[a[j]] = 1
                sumSoFar += a[j]
                if len(eleToFreq) == k:  # Updated condition
                    toReturn = max(toReturn, sumSoFar)
                
                if eleToFreq[a[i]] > 1:  # correction of condition
                    eleToFreq[a[i]] -= 1
                else:
                    eleToFreq.pop(a[i])
                sumSoFar -= a[i]
                i+=1
                j+=1
            else:
                if a[j] in eleToFreq:
                    eleToFreq[a[j]] += 1
                else:
                    eleToFreq[a[j]] = 1
                sumSoFar += a[j]
                j+=1

        return toReturn

Remarks on your code

There is some duplication of code in the if and else blocks which you can easily avoid by moving it out of there.

As you increment j consistently in each iteration, it is more pythonic to use a for..in loop, and in this case it is useful to use enumerate so that you don't need N either.

Instead of tuple assignment for initialising several variables to 0, you could chain the assignment and use a single 0.

It is not necessary to deal separately with the case where the number of unique items in the input is less than k. The main loop will deal with this, and so you avoid the work of creating a set of the input.

That leads to this code:

class Solution:
    def maximumSubarraySum(self, a, k):
        sumSoFar = toReturn = i = 0
        eleToFreq = dict()

        for j, value in enumerate(a):
            if value in eleToFreq:
                eleToFreq[value] += 1
            else:
                eleToFreq[value] = 1
            sumSoFar += value
            if j >= k - 1:
                value = a[i]
                if len(eleToFreq) == k:
                    toReturn = max(toReturn, sumSoFar)
                if eleToFreq[value] > 1:
                    eleToFreq[value] -= 1
                else:
                    eleToFreq.pop(value)
                sumSoFar -= value
                i += 1
        
        return toReturn

You could reduce the code more by using defaultdict, but it will not improve the performance.