Subarray Sums Divisible by K - Leetcode 974 - Python

Key Takeaways

hey everyone welcome back and let's write some more neat code today so today let's solve the problem subarray sums divisible by K I want to quickly mention before you even attempt this problem I would highly recommend solving leak code 560 I do have a pretty decent video on that problem this problem is a slightly harder variation of it it uses the core ideas this one just involves a little bit more math so the idea is given an integer array of nums uh let's take a look at the first example given we have an integer array with six numbers and we're given k equal to five so that's going to be important idea is we want to count the number of non-empty subarrays that have a sum divisible by K so brute forcing it how many subarrays does an array have roughly n^ 2 proportional to to n^ 2 we have n subarrays starting with the first element we have n minus1 subarray starting with the second element etc etc we could compute the sum of every single subarray and then we could check if it's divisible by K by modding it by K and if the result is equal to zero then we would count this if it's not we don't count it now after totaling up the count we return the account so pretty simple right can we really do better than an N squ solution how is that possible it's because we don't actually have to build the subarrays themselves we're counting the subarrays if we had to build every subarray there's going to be roughly N squared possibly like in the worst case so we wouldn't be able to do better than that but this time we're counting we don't have to build them therefore it might be possible that we can actually optimize this solution and we can indeed do that again we're going to use the core ideas from leak code problem 560 but even if you didn't know that let me kind of give you a bit of the intuition hypothetically if there was a way to optimize it this is what the train of thought would be instead of building every subray take a shortcut hypothetically do this maybe let's consider all subarrays starting at this element that would be one possible way to do it or we could say let's consider all subarrays ending at this element and then like iterate over the input it's kind of weird to do it starting at this element like how do we if we're going from left to right suppose and we're trying to get all subay sums divisible by K starting at this element kind of hard to do that if we haven't even looked at the other elements yet right so the train of thought is at this element let's get all subay sums ending at this element it's pretty easy to do for the first case with just one element we have the sum itself it's just four is this divisible by K which in this case is five well we mod it and we don't get zero we get a remainder of four so it's not we don't count it to our result our result is still zero now let's try the second thing now we have this our total sum so far is nine and we want to ask the question what are all subay sums count all subay sums ending at this element and check if they're divisible by K well there's two cases there's nine a mod that by five and we get a remainder of four so that's not it we could also try five by itself right all subs ending here include this one five itself so five moded by five gives us a remainder of zero so in this case obviously we add one to the result the result is going to be one now so what's the pattern well let's try one more time here well first of all the sum here is nine once again we try this itself all subray sums ending at this element 9 moded by 5 is going to be four that doesn't work okay try this subay sum which is just five five modded by five is indeed zero we counted one more now try this subay sum ending at this element just itself removing these two we get zero modded by five which is again zero so this one actually counts like you can see even in the output here they say the subay with just zero is technically divisible by K so that does count the math works out great so so far we counted two subay sums ending at this element and we counted one subay sum ending at this element do you notice the pattern well to get potentially all subay sums ending at this element we try the whole thing we try removing this and then we try removing this obviously if we do it manually it's going to still be an N squ solution is there a shortcut that you notice through all of that the idea is that for the math to work out if we're keeping track of the total sum every time we get to a new element we're keeping track of the total sum from the beginning up until this point we're trying to count the sub Rays ending at this element we have our let's say current sum or total sum we want it modded by K to be equal to zero but we're also allowed to have that sum minus a prefix of the array and that modded by K can also equal zero now the hard part is that these prefixes could be this this or this or this or whatever right any prefix so obviously we'd have to Loop to do that right is there a better way of course there is it's not easy to come up with but why not solve the mathematical formula for this what does the prefix have to be such that this equation is true well we kind of saw an example of it earlier let's just look at the example let's say the sum here is nine if we mod this by five we get a remainder of four we want a remainder of zero but we got a remainder of four so obviously we're trying to remove a remainder of four from the beginning right a prefix of four so the remainder that we got is exactly what we're trying to remove but it actually works in another case as well if we take the entire prefix 9 itself and take 9 - 9 and we remove that we also then will get a remainder of zero if we take 9 - 4 we get a remainder of zero and if we do 9 - 9 we get a remainder of zero so what's the pattern here what is exactly the prefix that we're trying to remove it doesn't necessarily have to be the remainder it's just that the prefix itself needs to have the same remainder so for example we have 9 here 9 moded by 5 which is K is going to be four we want to remove any prefix that is equal to four or also has a remainder of four so a prefix of nine is also what we're trying to remove because nine would also have a remainder of four after modding it by K so that's the math I know it's not super intuitive especially if you haven't solved leak code 560 but if you have this is a somewhat reasonable problem if you deeply understand that this one just involves a bit more math here to summarize everything we're going to be counting the prefixes but not based on their sum but we're going to be counting the prefixes based on their remainder so here on the left I'm going to have a happy little hashmap and it's going to be called count you could call it whatever you want we're going to have the remainder as the key and the value is going to be the count initially we're going to actually have a default value in here we're going to say there is a prefix that has a remainder of zero and a count of one and the reason for that is it'll probably make more sense in the code or I might explain it in a few minutes but for now just kind of leave this here we're going to keep track of what the current sum is starting at the first element we have four we're going to get the remainder of this modding it by K we're going to have a remainder of four well that's not good we want all the remainders to be zero this is not so now we want to ask is there a prefix from here that we can remove such that the remainder of this would actually be zero well for us to do that we have to have a prefix with a remainder of four we don't have any of those yet so we can't create a non-empty subarray okay that's fine but now we're going to say we've counted one subarray one prefix that had a remainder of four so we're going to add the count here remainder of four has a count of one now we're going to look at this sum we're trying to count all subarray sums ending at this element right now our sum is nine mod this by five we get a remainder of four it's not zero so we check how many subay sums did we have with a remainder of four because that tells us how many subarray sums we can create here that are going to be divisible by K in this case we had exactly one that tells us we can remove that one subarray and then the remainder of this would be equal to zero so we counted one so now the result is going to be set to one and we're going to take this and update our count hashmap we're not going to say nine is equal to one we're going to increment this one nine was a sum a prefix sum it had a remainder of four therefore we've counted two prefix sums with a remainder of four so the count now is going to be two now we're going to look at this the sum is nine we mod it by five we get a remainder of four so we check is that in the hashmap yes it is so now the result is going to be incremented by two those subarrays would look like this this and of course this there were two prefix sums we could remove such that the remainder of the leftover would be four or the leftover would be zero but those prefix sums would have a remainder of four and so once again we counted a prefix with a remainder of four so this is going to be incremented by one so now we've counted three prefix sums with a remainder of four so now we look at this subarray total sum is seven mod that by five we get a remainder of two how many subay sums have we seen with a remainder of two I don't see any here we don't add anything to our result our result I believe at this point is going to be three we don't do anything with the result but now we've counted one prefix sum that had a remainder of two so we can update our hashmap with that information there's one with a remainder of two so now I'm adding this -3 I believe we have four so four modded by 5 is going to be equal to four how many prefix sums did we have with remainder of four we counted three exactly what would those subarrays look like well this has a subarray sum that's divisible by K getting rid of this this also has a subray sum divisible by K5 is technically divisible by 5 and then this one as well this is A5 it's divisible by five so we've counted those three so now we're going to add three to our result so right now our result would be six and since we counted one more prefix sign that is having a remainder of four we're actually going to set this now equal to four we're going to increment it by one now lastly here we have a total sum of five and we take this and we divide it by K we get zero but the way we're updating our result is not based on this it's based on how many other prefix sums that we have with the same remainder and so we counted a single one and you can think of that as like the empty prefix um technically with an empty prefix we remove that empty prefix we have this array left over this technically is divisible by K so now you can kind of see why we initialize it like this because we do still want to consider that case it's mainly just to make the math work out and to keep our code relatively clean there are alternatives to this but if uh this didn't make sense imagine something like this imagine we had an array that just has five we take five modded by K which is zero and we'd want to see that already we have counted a prefix sum with a remainder of zero which is one so we'd add one to our result that's the idea here in summary that last one was a single subarray sum so we had 6 + 1 the result is seven and I believe that is the result that we had and if I go backwards yep you can see here we had seven as the expected output so that's how we're going to code this up obviously it's a linear time solution and you might think that the hashmap also is going to be linear time in the worst case but actually it's a bit more efficient than that because think about how many possible remainders we could have when we're modding by five the remainder could only be either 0 1 2 3 or four right 0 through K minus one so the max size of the hash map is actually going to be K so this is the time complexity this is the spe base complexity okay so now coding it up um we're going to keep track of the current sum I guess I'll call it the prefix sum just to be very clear initially it's going to be zero and we're also going to have our result that's what we're going to return that's the count of the subay sums divisible by K and also we're going to have that hash map I'm going to call it prefix count I guess but there might be a more descriptive name for it I'm actually going to set it to a default dictionary just cuz that's always ni nice and I'm going to initialize it with a single pair mapping zero as the remainder to a count of one so then we're going to go through every number in the input we're always going to add it to the prefix sum then we want to compute the remainder so let's do that remainder is going to be prefix sum modded by K now we want to know how many prefix sums were having a remainder of exactly this because with those we're going to add them to the result so you could do it like this if remainder is in prefix count then add to the result the amount that we counted so prefix count with this remainder add that to the result and also after we're done with that let's update the prefix count by adding one more of the current remainder to it so this is exactly what I talked about in the drawing explanation now it's not easy to come up with I admit but this is the code we could make it slightly more concise with a default dictionary if a key doesn't exist this will evaluate to zero we won't get like a key out of bounds error or key doesn't exist error so we can actually move this outside of the condition and get rid of the condition if you really want to it's not like a major optimization or anything but that's kind of the power of a default dictionary so this is is the code let's run it to make sure that it works and as you can see on the left yes it does it's about as efficient as we can get it if you found this helpful check out NE code. thanks for watching and I'll see you soon