Understanding Count Triplets HackerRank

I have been working on this challenge: Count Triplets, and after a lot of hard work, my algorithm did not work out for every test case.

Since in the discussion, I have seen a code and tried to find out the real functionality of the code, I am still not able to understand how, this code works.

Solution:

from collections import defaultdict

arr = [1,3,9,9,27,81]
r = 3
v2 = defaultdict(int)
v3 = defaultdict(int)
count = 0
for k in arr:
    count += v3[k]
    v3[k*r] += v2[k]
    v2[k*r] += 1
print(count)

The above code works for every test case perfectly. I have tested for value of k, v2, v3 to understand but still don't understand how the code works so smooth with the counting triplets. I cannot think of that solution in my dreams too. I wonder how people are so smart to work out this solution. Nevertheless, I would be glad if I would get the proper explanation. Thanks

Output for k,v2,v3

from collections import defaultdict

arr = [1,3,9,9,27,81]
r = 3
v2 = defaultdict(int)
v3 = defaultdict(int)
count = 0
for k in arr:
    count += v3[k]
    v3[k*r] += v2[k]
    v2[k*r] += 1
    print(k, count, v2, v3)

OUTPUT

1 0 defaultdict(<class 'int'>, {1: 0, 3: 1}) defaultdict(<class 'int'>, {1: 0, 3: 0})                                                  
3 0 defaultdict(<class 'int'>, {1: 0, 3: 1, 9: 1}) defaultdict(<class 'int'>, {1: 0, 3: 0, 9: 1})                                      
9 1 defaultdict(<class 'int'>, {27: 1, 1: 0, 3: 1, 9: 1}) defaultdict(<class 'int'>, {27: 1, 1: 0, 3: 0, 9: 1})                        
9 2 defaultdict(<class 'int'>, {27: 2, 1: 0, 3: 1, 9: 1}) defaultdict(<class 'int'>, {27: 2, 1: 0, 3: 0, 9: 1})                        
27 4 defaultdict(<class 'int'>, {27: 2, 1: 0, 3: 1, 81: 1, 9: 1}) defaultdict(<class 'int'>, {27: 2, 1: 0, 3: 0, 81: 2, 9: 1})         
81 6 defaultdict(<class 'int'>, {1: 0, 3: 1, 243: 1, 81: 1, 9: 1, 27: 2}) defaultdict(<class 'int'>, {1: 0, 3: 0, 243: 1, 81: 2, 9: 1, 
27: 2})

Solution 1:

1. The problem

The function has two parameters, namely:

  • arr: an array of integers
  • r: an integer, the common ratio

So, the input can be something like

arr: [1, 2, 2, 4]
r: 2

The goal is to return the count of triplets that form a geometric progression.


2. How to solve it

To solve it there's variou ways. For instances, from SagunB based on the comment from RobertsN

  • Can be done in O(n) -> single pass through data
  • No division necessary and single multiplications by R are all that's needed
  • Using map(C++) or dict(Java, Python) is a must -> can be unordered map (saves O(logN))
  • Try to think forward when reading a value -> will this value form part of a triplet later?
  • No need to consider (R == 1) as a corner case
from collections import Counter

# Complete the countTriplets function below.
def countTriplets(arr, r):
    r2 = Counter()
    r3 = Counter()
    count = 0
    
    for v in arr:
        if v in r3:
            count += r3[v]
        
        if v in r2:
            r3[v*r] += r2[v]
        
        r2[v*r] += 1

    return count

Or like you said

from collections import defaultdict

# Complete the countTriplets function below.
def countTriplets(arr, r):
    v2 = defaultdict(int)
    v3 = defaultdict(int)
    count = 0
    for k in arr:
        count += v3[k]
        v3[k*r] += v2[k]
        v2[k*r] += 1
    return count

3. End result

Both cases will pass all the current 13 Test cases in HackerRank

It works


4. Explanation of your case

Comments from RobertsN pretty much explain your code (which is very similar to yours). Still, for a better clarification to understand how the code works, just print the what happens to count, v2 and v3.

Assuming you'll have as input

4 2
1 2 2 4

The expected output is

2

Also, we know that by definition both v2 and v3 will look like

defaultdict(<class 'int'>, {})

which leaves the for loop left to understand. What can cause some confusion there is the operator += but that was already addressed by me in another answer.

So, now to understand the rest we can change the loop to

for k in arr:
    print(f"Looping...")
    print(f"k: {k}")
    print(f"v3_before_count: {v3}")
    count += v3[k]
    print(f"count: {count}")
    print(f"k*r: {k*r}")
    print(f"v3_before: {v3}")
    v3[k*r] += v2[k]
    print(f"v3[k*r]: {v3[k*r]}")
    print(f"v2[k]: {v2[k]}")
    print(f"v3_after: {v3}")
    print(f"v2_before: {v2}")
    v2[k*r] += 1
    print(f"v2_after: {v2}")
    print(f"v2[k*r]: {v2[k*r]}")

Will allow you to see

Looping...
k: 1
v3_before_count: defaultdict(<class 'int'>, {})
count: 0
k*r: 2
v3_before: defaultdict(<class 'int'>, {1: 0})
v2_before_v3: defaultdict(<class 'int'>, {1: 0})
v3[k*r]: 0
v2[k]: 0
v3_after: defaultdict(<class 'int'>, {1: 0, 2: 0})
v2_before: defaultdict(<class 'int'>, {1: 0})
v2_after: defaultdict(<class 'int'>, {1: 0, 2: 1})
v2[k*r]: 1
Looping...
k: 2
v3_before_count: defaultdict(<class 'int'>, {1: 0, 2: 0})
count: 0
k*r: 4
v3_before: defaultdict(<class 'int'>, {1: 0, 2: 0})
v2_before_v3: defaultdict(<class 'int'>, {1: 0, 2: 0})
v3[k*r]: 1
v2[k]: 1
v3_after: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 1})
v2_before: defaultdict(<class 'int'>, {1: 0, 2: 1})
v2_after: defaultdict(<class 'int'>, {1: 0, 2: 1, 4: 1})
v2[k*r]: 1
Looping...
k: 2
v3_before_count: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 1})
count: 0
k*r: 4
v3_before: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 1})
v2_before_v3: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 1})
v3[k*r]: 2
v2[k]: 1
v3_after: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 2})
v2_before: defaultdict(<class 'int'>, {1: 0, 2: 1, 4: 1})
v2_after: defaultdict(<class 'int'>, {1: 0, 2: 1, 4: 2})
v2[k*r]: 2
Looping...
k: 4
v3_before_count: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 2})
count: 2
k*r: 8
v3_before: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 2})
v2_before_v3: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 2})
v3[k*r]: 2
v2[k]: 2
v3_after: defaultdict(<class 'int'>, {1: 0, 2: 0, 4: 2, 8: 2})
v2_before: defaultdict(<class 'int'>, {1: 0, 2: 1, 4: 2})
v2_after: defaultdict(<class 'int'>, {1: 0, 2: 1, 4: 2, 8: 1})
v2[k*r]: 1

and extract the desired illations. What can we observe from that?

  • count increases in the last loop from 0 to 2.
  • k goes through all values of the arr - so it'll be 1, 2, 2 and 4.
  • in the initial loop, v3_before_count is {} and v3_before is {1:0}
  • etc.

Most likely this process will lead to questions and answering them will leave you closer to understand it.

Solution 2:

So the code is tracking potential pairs and triplets as it walks through the array.

For each value in the array:
    // Increment count by the number of triplets that end with k
    count += v3[k]
    // Increment the number of potential triplets that will end with k*r
    v3[k*r] += v2[k]
    // Increment the number of potential pairs that end with k*r
    v2[k*r] += 1  

The number of triplets for any given k is the number of pairs for any given k/r that we've encountered up to this point. Note throughout the loop, v3[k] and v2[k] will often be zero, until they hit our predicted k*r value from a previous iteration.

Solution 3:

I have been trying to make sense of it and finally, this C# code should be clear to follow

static long countTriplets(List<long> arr, long r)
{
    //number of times we encounter key*r
    var doubles = new Dictionary<long, long>();
    //number of times we encounter a triplet
    var triplets = new Dictionary<long, long>();
    long count = 0;
    foreach (var key in arr)
    {
        long keyXr = key * r;

        if (triplets.ContainsKey(key))
            count += triplets[key];

        if (doubles.ContainsKey(key))
        {
            if (triplets.ContainsKey(keyXr))
                triplets[keyXr] += doubles[key];
            else
                triplets.Add(keyXr, doubles[key]);
        }

        if (doubles.ContainsKey(keyXr))
            doubles[keyXr]++;
        else
            doubles.Add(keyXr, 1);
    }
    return count;
}