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.
