Count Good Triplets - Leetcode 1534 - Python

Key Takeaways

Hey everyone, welcome back and let's write some more neat code today. So today, let's solve the problem count good triplets. So the problem description is pretty straightforward in terms of implementing the brute force solution. I'll go ahead and read the problem quickly and implement that solution, but I want to spend most of the video focusing on the more optimal solution, which is going to require us to actually make some observations. You can't just blindly read the problem and then expect to arrive at the solution. You actually have to think about it. I like to use the word making observations about the problem, figuring out what the tricks are. Now, for an easy problem, I will say that that more optimal solution is definitely difficult. It's not an easy solution. It's close or I would definitely just say it is a medium solution. So, don't feel bad if you weren't able to come up with that. But anyways, let's get into it. The idea here is that we are given an array of integers. We want to get triplets from that list. Now, we can't use the exact same element multiple times. But if suppose a number shows up multiple times, then we can use the same number, but we can't use the same element if that makes sense. So like for one here, there are two of them. We could use both of them. We want to pick three numbers. So I'm going to draw them like this. Three slots. Even though the index well well I guess I'll just draw it over here. So like this number is going to be at index i. This number is going to be at index j and this number is going to be at index k and they need to be at different indexes obviously. But also this needs to be true. This needs to be at a greater index than that one and this needs to be at a greater index than that one. Whatever numbers we pick here, let me just arbitrarily call them x, y, and z. As long as they satisfy this condition and they satisfy these three conditions over here, that's where most of the complexity is, these three conditions, then we can say that this is a good triplet. And what we want to do is count all of the good triplets. So the good thing is that this condition is very trivial to check. You don't really even have to turn your brain on. All you have to do is recognize that okay, this is a math formula. the number at index i, subtract from it the number at index j and get the absolute value of it. So in that case, it doesn't really matter how we do it. We could do x - y or we could do y - x because we're taking the absolute value of it anyway. When you take the difference between two numbers, think of it like this. You have two numbers on like a number line. When you take the difference between them, you're just getting the distance between them. If you take the absolute value, then you're just going to get that as a positive number. That is pretty straightforward. And the other two conditions are pretty much the same. So J and K. So that means Y minus Z the absolute value of that. And then the other one is the absolute value of this minus this. And specifically for each of those we want to know that uh the first one is less than or equal to the value a which is provided to us as a parameter. The second one less than or equal to the value b. The third one less than or equal to the value c. Okay. So there's a lot going on, but the main thing to recognize for the brute force solution is that we just need triple nested loops. So like just having three nested loops and then inside of the last loop, we'll just check this condition which will allow us to determine if it's a good triplet or not. The nested loops are basically to just enumerate every single possible triplet and then we'll use an if statement to check that it's actually true or not. If it is, we can increment our count. Otherwise, we do not increment the count. So this will be a solution which is bigo O of N cubed because we have triple nested loops. Let me quickly code this up. I'm going to go kind of fast because this should be mostly straightforward. We have our result that's going to be zero. That's what we're going to return. I'm going to have three loops. I'm going to have my I pointer go from the beginning all the way up until the end. And actually I'm going to use a variable to get the length of the array because we're going to need that a few times. We could go up until the end of the array, but we know that I is the first element. We want at least two elements after it. So to make it easier, I'm just going to put minus2 here. You could get away with not doing that, but I'm going to do it anyways because inside of here, we're going to have the second loop J. And that's not going to start from the beginning of the array. It's actually going to start right after I because of course this element has to be after it. So this is going to be I + 1. I'm going up until the end minus one because we need one slot left for the last element which is K in range and this one will start right after J + 1 up until the end. The reason I say we don't actually need this and this is because I mean here we're starting at I + 1 anyways and if we were at the last position I + 1 would have been out of bounds. This loop wouldn't have executed anyway. So uh I'm just putting these here to make it more clear but you don't need them. Try to notice how even though this is a brute force solution, clearly I have a very deep understanding of what's going on here. So maybe if you're a beginner, you were able to implement this, but maybe you don't have quite as deep of an understanding as I do. That's perfectly fine. That's why we're here to practice. Anyways, now the last thing to do is just the condition. And it's just going to be a pretty fat condition. And then inside of there, if it executes to be true, we will increment our count. So in here I'm going to put three conditions for our if statement. The first one is going to be absolute value of the number at i minus number at j. That has to be less than or equal to a. And now I'm just going to copy and paste this cuz it's going to be pretty repetitive. This one is going to be less than or equal to b. This one less than or equal to c. And then I think for this one it's going to be j minus k. And again like the order of these doesn't matter. I could have put k here and I could have put j over there. Um, and lastly, I minus array at K. So, that's pretty much it. Pretty trivial solution. Let's run it. And you can see it works. It's pretty efficient, but they're definitely is a more efficient one. That's what I'm going to show you. And this brute force solution in some ways is definitely related to that other solution. Let me kind of give you a brief drawing explanation. So, how do you even approach a problem like this one? Because at the end of the day, what we're doing is counting triplets. I mean, how can you possibly count triplets without looking at every single triplet? So obviously, if there is a solution to this, it's going to involve taking our result and not just incrementing it by one, but incrementing it by a larger number. So think of our triple nested loops. How could we maybe get rid of one of those loops? We have our 4 I. We have then inside of that the 4 J. And then lastly, we have the 4k. So the first thing that you might try to do is figure out if you can eliminate the 4k loop. Just by fixing these two pointers, can we somehow know in advance what's going on with the rest of the array? Well, before I go too down this rabbit hole, I want to mention that the main reason that we can solve this more optimally is the fact that there's a constraint on the input. I believe it's that the values in the array are going to be between zero and I think up to a th00and. They put that somewhere like down in the description. So that's the reason that we can solve this more optimally. I always think it's kind of stupid when like the constraints allow us to do that, but practically speaking that is a good reason to optimize the problem. Like if the constraints were actually like this now, how many problems like in real world would the constraints be that small? Who knows? But anyway, so that's going to be the idea here. And specifically for this, we can eliminate one of the loops. I'm actually not going to eliminate this loop. It might be possible to do that, but it's a lot easier to eliminate the outer loop. And I'll show you how we do that. It's going to involve some counting. Kind of related to how prefix sums work. It's kind of going to be prefix counting. So, this is definitely an advanced concept, I will say, though. So, if you're a beginner, don't expect to understand this stuff. I would highly recommend checking out neat code.io. I believe prefix sums are covered in the advanced algorithms course if you're interested in that, but there's a lot of other free resources on there as well. And remember that what we're ultimately trying to do is count when the conditions are specifically met. So again, let's go back to a simple picture so that we can understand. Think of it in terms of like the number line or something like that. We already know like the relationship between the indexes, but let's suppose that we fixed these two loops. And so my current number at index I, let's say it's over here. Uh let's just put I over there. And let's say my number that's currently at index J on the number line is something like this. What we want to confirm is that first of all, the distance between these two numbers, not the distance between the indexes. So just want to make that clear. The actual numbers themselves the distance between those is going to be less than or equal to a which in this example is going to be seven. So let's say that that is true. Okay, that's true. Now what I want to know and we started from the beginning of the array remember that. So let's say like we've processed some portion of the array. Some portion of the array is left over. I want to know now how many of the rest of these values satisfy the other two conditions. Well, let me draw out the other two conditions on this number line for a second. For J over here, the difference between it and the number at index K has to be less than or equal to B. So, let's say I I just draw that over here. Like I draw a little B segment over here. Let's say the distance of this is B and the distance of this is also B. So any number that falls within this range plus or minus B from J is going to be valid for this second condition. Perfectly fine. And for I, it's similar down here except the value this time is going to be C. So I'm just going to, you know, draw some random line. Let's say this is of distance C. And let's say over here as well. I'm going to draw it kind of underneath this one. Let's say this is of distance uh C as well. So now you tell me for a second what were we even trying to do again? Well, remember we were trying to get rid of this last loop fork because I don't want to have to go through all of those elements to know which one of them are valid. So I did a little shortcut over here. This picture tells me that okay this green segment satisfies the second condition and the orange segment satisfies the third condition. So basically I can pick a k value that is here and here. Now what if I put my k value over here? Just tell me for a second if I put my k value over here you tell me for a second is that valid or is that not valid? Well it definitely satisfies the green condition. The second one it does not satisfy the orange condition though. Remember all of the conditions have to be satisfied. So once you realize that you realize that it's the intersection that we care about over here. So in other words, we would want to take this segment. And how exactly do we get this segment? Well, in terms of the code, that's actually going to be pretty easy. Intuitively, it might not be easy like off the top of your head, but it's just going to be taking the minimum of the two end points, like the two right values. Take the minimum of those. And for the left points, take the maximum. In other words, for these two end points, we want to get the rightmost one because again, that's where the intersection is going to happen. And for these two points, we want to get the leftmost one because we're trying to get the range that both of them satisfy. So, okay, once you get this range, okay, now we have it. We have that range. Let's just say it's from X all the way to Y. So now we want to know how many values in the remaining portion of the array fall within this range. For us to do that, we're going to now have to scan through those elements. So it didn't really help us, did it? Like we're still going to need another loop down here. That's why instead of eliminating the K loop, we actually eliminate the first loop because now imagine if we did it the opposite way. Imagine uh these two loops are good, but we got rid of this loop. Instead of asking, okay, from the remaining portion of the array, how many elements fall within that specific range, we're going to instead ask ourselves how many previous elements that we've already seen fall within that range. In other words, as we're scanning through this, we're not only going to be doing this stuff, but we're going to be doing some counting because when we're doing that counting, we're going to be able to recognize how many previous elements fell within that particular range. And the counting is going to be done with an array that's going to be of dimensions. I think from zero all the way like the size of it is going to be uh 1,000. Well, I think 1,0001 because we have zero and then the max value could be a th00and something like that. But I think this part is probably going to be just easier to explain in the code. The idea is going to be prefix counting. A concept that honestly doesn't come up too often. So, you probably definitely want to have a deep understanding of prefix arrays to solve this follow-up solution that I'm going to show you. Anyways, I'll analyze the time complexity in the code as well, but it's mostly going to be I think n squared plus n * the max value in the array, which let's call it m and m could be up to a,000 in this case. So, that's why this is more efficient. And I think with the constraints of the problem, n was actually up to 100. So if you take like 100 squared and then this which is going to be 100 * 1,000 I think this is going to be the bottleneck but it's still better than the previous solution which would have been 100 cubed. I think that would have been like a million. This is 100,000. So it's better by like a factor of 10 at least in the worst case with these specific constraints. Okay. So I'm actually going to leave most of this code here. I'll zoom in just a tad bit. And I'm going to go ahead now and just get rid of this outer loop. And I'll show you what I'm going to do with it in just a second. And I'm going to um have my J not start at I + 1, but it's going to start from the beginning. Now, you might think, well, if J starts from the beginning, then we won't have a valid I. And that's fine. For the first iteration, we actually won't need a valid I. We'll show you, or I'll show you what goes in over here. A K will do the exact same thing, though. And this is going to be solved by this specific data structure. It's going to be an array. I call it prefix count. It's initially going to be all zeros and of size 1,01. And what this is going to do is calculate the prefix counts. So specifically prefix count of some number x is going to map to the count of the numbers that are less than or equal to x. So in other words like the upper bound like X is the upper bound and then it'll tell me how many numbers are less than or equal to X. How many have we seen? That's where the counting stuff is going to be done more efficiently. Okay. So now we do this. We have our J and we have our K are fixed. What are we going to do? Let's just assume that this data structure is populated correctly. What exactly are we going to do? Well, what we're specifically looking for is this. How many values before J where the absolute conditions are met? And remember those absolute conditions are like those intervals that I was showing you. So think of it this way. We have our number at index j. Now we're actually looking for the i value. And remember the difference between i and j was determined by a over here. So we can say our current value plus a. And we know that a is always going to be positive. Similarly, we could do our number at K plus I think it was C with I. So I and K. The difference between those has to be less than or equal to C. So here I'm going to do plus C. And then I'm going to copy uh this line down and I'm going to do the opposite where I'm going to take the difference. So I'm going to subtract from it. So array of K minus C and array of J minus A. So we want to look for the intersection. So among these two, I don't want the bigger one. I actually want the smaller one. So I'm going to take the minimum of these two and that's going to determine my right boundary. So I'm going to call that r because we know that a and c are positive. So by adding to these, we're only making them bigger, but we want the minimum of the two. And then over here, we're going to do the same thing except kind of take the maximum cuz we're making both of these smaller. And the bigger one is going to be where the intersection is happening. Actually, now I I I remember we forgot to do the other thing where we never compared the values at J and K. So let's uh remember to do that. So before we even do these two down here, we have to make sure that these two are close enough. So that can be done pretty easily. Absolute of J minus array of K. And that has to be less than or equal to I believe B because we already used A and C down there. So okay. So this is where we're doing everything. And now whoops that we have the boundaries of our intersection. What we would want to do is this. We would want to say assuming our prefix count is correct. Prefix count of r is going to tell us how many numbers are less than or equal to this. But we don't want that exactly. We want within this entire range between r and l. So from that we will actually subtract prefix count not of L but L minus one because we want the range inclusive. So this is just kind of how prefix sums work. If you're not familiar with this line of code, believe it or not, this line is actually the easy part. This stuff is hard. This stuff is hard, but this line is actually the easy part. So if you're not familiar with prefix counting or prefix sums in general, you might just need a little bit more practice, and that's perfectly fine. One edge case here is that okay well what if prefix of L minus one is out of bounds like what if L was zero well then this would give us an index out of bounds error we don't want to deal with that there are multiple ways to deal with it we could have done like an offset over here we could have made this of size 1002 sorry then we could have done some offsetting over here but I'm not going to do that the other way to handle this is sort of with a turnary operator so I'm going to actually wrap this in a turnary I'm going to say this is actually going to be zero if L is equal to zero. Otherwise, go ahead and index the array because if L is zero, we don't want to index the array. Only if it's not zero. So, there's that. And then this is the count. This is how many numbers came before index J that met the absolute conditions. And so, this is what we would want to add to our result. Now, there's a couple bugs and again, this is the hard part. First of all, L could technically be out of bounds. Like if we took something and subtracted from it, it might have gone too small. So if it's less than zero, we want to reset it. So we could say L is the max of itself and zero. So if it was less than zero, it would be updated. Same with R. We don't want it to get too big. We don't want it to be bigger than a th00and cuz our prefix count array isn't bigger than a th00and and we don't have any values in the array that are bigger than a th00and. So we can say it's going to be the minimum of itself as well as a th00and almost done. But there is a possibility that actually our two pointers might have crossed each other. So here we need to validate that left is less than or equal to right only then we want to do this because otherwise we would possibly get a negative count that is added to the result which doesn't really make sense. And an example where that could occur is maybe like array of J. I I think the example would look something like this where we have like let's say a segment and then another segment that are not overlapping each other and then so if you try to take like the intersection of those it doesn't really make sense and so like you get like these two right values and you get the minimum of those which is going to be this one and then you get the two left end points and you get the maximum of those two which is going to be this one. So okay so this is or sorry it's going to be this one. So this is left this was right. It doesn't really make sense for our left pointer to be here and our right pointer to be here. So that's kind of the example where this would sort of break down. That's why we add this guard over here. Again, this is not easy, is it? Even though this problem is marked as an easy, don't let leak code fool you. Okay, lastly, how exactly are we actually going to have the prefix count be updated? Because we assumed here that it's going to work, but when do we actually update it? Well, after we're done processing every possible K value here, then before we go to the next iteration of this loop, go to the next J value, we want to say, okay, everything that we saw before J should be in the prefix count. So, we're going to say, okay, the current value that we just saw should be added to prefix count. Well, that's going to require a loop. We can't just do something like this. We can't just do prefix count of array at J + 1. That is necessary, but that's not alone sufficient because our prefix count has to update every index greater than this one as well because that's what our prefix count is storing. All of the numbers less than or equal to x. So if I saw the number four, for example, I would want to set that index to be one, but I'd want to set every index greater than that as well to be one. So I'm going to have a loop. I'm going to say okay my index is going to go in the range starting from J all the way up until the last value which is going to be a th00and so we put a 101 and then I'm going to take this line cut it and I'm going to move it uh down here and then I'm going to swap this with index. Now I hope it's obvious that this solution only works because this constraint is pretty small. So let's just zoom out see the entire code. So this is pretty much it. I'll go ahead and give it a run. And you can see here it works. And we finally reached that second peak over here. If you found this helpful, definitely check out necode.io for more. Thanks for watching and I'll see you