Good Subarrays with Python and Javascript

With Python and Javascript solutions: Let's learn how to think about Good Subarrays, a 1600-point codeforces problem

Understanding the Problem

Good Subarrays is a 1600-point problem at codeforces, I hope you have read the problem carefully and have tried solving it. This problem basically states as its name sounds to say, to get a good subarray. But what is good subarray?

Per each test case, you’re given a string of numbers and the length of this string, and what’s required is to get the good subarrays. A good subarray is a subarray whose length equals to the sum of its elements.

How to Think for a Solution

This problem has three inputs:

  • input for test cases
  • input for the length of the array
  • input for the string of numbers per each test case

You can first think of this problem in a brute force way to loop over the string to get each subarray possible. For example, for the ‘120’ case you loop from 1 and you can detect the first good subarray which is ‘1’ having a subarray length of 1 and a prefix sum of 1, and then you proceed with the first loop from 1 until you get 2 and 0 so you can get another good subarray which is ‘120’ which has a length of 3 and a prefix sum of 3, and then loop again from 2 until you get 0 so you can get one more good subarray which is ‘20’ having a length of 2 and a prefix sum of 2.

But this way will get you in trouble when the number of elements is big. For example, what if you have 10,000 elements .. then you have to loop 10⁸ times! That’s not practical, what we need to do is to try to memorize the number of good subarrays. We can do that by doing hash maps.

The keys for the hash map (hash) are the difference between prefix sum (s) of that subarray and length of that subarray (l) and we need to loop just one time so that the time complexity for this solution will be O(n) instead of the brute force way which was O(n²). The values for that hash is the old values stored for the keys (hash[s — l] ++). At each iteration, we need to calculate how many good subarrays we have (g) and how many the prefix sum is accumulated (s) and the length of the subarrays (l), and then we need to update the hash map.

For the same example ‘120’, before we iterate:

  1. We get a total sum of zero because we didn’t actually sum any number yet so s = 0
  2. The length of subarray we have is zero because we don’t have any subarray yet so l = 0
  3. The difference between prefix sum and length of the subarray is zero,
    so s-l= 0
  4. That will give us no good subarray as well, so g = 0
  5. The value of the hash map key (s — l which is 0) is 1 because we can always find subarray that has prefix sum of 0 and subarray length of 0 if I don’t take any element so hash[0] = 1

For the first iteration ‘1’:

  1. s = 1
  2. l = 1
  3. s-l = 0
  4. A good subarray is the sum of the previous number of good subarrays (g) and the value stored in the hash map for that specific key (s-l), so g now equals 0 (g) + hash[0] (1) = 1
  5. hash[0] = 1 (previous hash[0] + 1) = 2

For the second iteration ‘12’:

  1. s = 3
  2. l = 2
  3. s-l = 1
  4. g = 1 + hash[1] (0) = 1
  5. hash[1] = 1

For the third and last iteration ‘120’:

  1. s = 3
  2. l = 3
  3. s-l = 0
  4. g = 1 + hash[0] (2) = 3 *good subarrays*
  5. hash[0] = 3

That’s how I thought about this problem and below is my implementations in Python and Javascript:



Motivated by

Other problem-solving stuff

If you want to see more problem-solving blog posts, check out:

Join the conversation

Get FREE coupons and discounts
on my upcoming courses
when you subscribe!

Thank you! Your submission has been received!
Oops! Something went wrong while submitting the form.