Understanding Count Triplets HackerRank

Question

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})

Answer 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

---

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.

Answer 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.

Answer 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;
}