The Problem

A link to the HackerRank problem: https://www.hackerrank.com/challenges/count-triplets-1

Problem Statement

You are given an array and you need to find number of triplets of indices  such that the elements at those indices are in geometric progression for a given common ratio  and  .

For example, . If , we have  and  at indices  and .

Input Format

The first line contains two space-separated integers and , the size of and the common ratio.
The next line contains space-separated integers .

Output Format

Return the count of triplets that form a geometric progression.

Constraints

• 
• 
• 

My Solution

Language: JavaScript

I generally like to start with the simplest, most brute-force method. Then I’ll try to be smart make more efficient algorithms. What I like about challenge sites like HackerRank is that the tests time out quickly, forcing you to write optimal code.

Iteration 1: simple brute force attempt

Here’s my first attempt at the problem

const solve = (nums, ratio) => {
  let count = 0
  
  for (let i = 0; i < nums.length - 2; i++) {    
    
    for (let j = i + 1; j < nums.length - 1; j++) {
      
      if (nums[j] / ratio !== nums[i]) continue
      
      for (let k = j + 1; k < nums.length; k++) {
        if (nums[k] / ratio !== nums[j]) continue
        count++
      }
    }
  }
  
  return count
}

The implementation is fairly simple. My approach was to do a triple nested loop through the number array. Every time it would check if the final number was next in the geometric sequence and then add to the count if it was.

As you might expect, this was not very fast. It’s which is…not impressive. That’s okay, though because I was planning on it failing. Benefits of low expectations.

I didn’t want to spend too much time doing a brute force solution but I wanted to see if I could cut the time down meaningfully at all. I ended up adding one thing.

const solve = (arr, ratio) => {
  let count = 0
  let nums = arr.slice().filter(num => num === 1 || !(num % ratio))
  
  for (let i = 0; i < nums.length - 2; i++) {    
    
    for (let j = i + 1; j < nums.length - 1; j++) {
      
      if (nums[j] / ratio !== nums[i]) continue
      
      for (let k = j + 1; k < nums.length; k++) {
        if (nums[k] / ratio !== nums[j]) continue
        count++
      }
    }
  }
  
  return count
}

I tried to cut down the number of items in the array before going down the black hole that was the triple for loops. Surprisingly it did help a bit: 2 more test cases passed. It was a very half-hearted measure though. I could do better.

Iteration 2: using dictionaries so only one loop is required

Time to write some real code! The most important factor is to only iterate through the array once. To do that, I saved a dictionary of pairs and possible pairs as I went along (this problem was in the ‘dictionary and hashmaps’ category so that was kind of a hint)

const solve = (nums, ratio) => {
  const groups = []
  let count = 0
  
  for (const num of nums) {
    if (num % ratio !== 0 && num !== 1) continue
    
    let inGroup = false
    for (const group of groups) {
      if (group.nums.length === 1) {
        if (num / group.nums[0] === ratio) {
          group.nums.push(num)
          group.counts[1]++
        }
        if (num === group.nums[0]) {
          group.counts[0]++
          inGroup = true
        }
      } else {
        if (num / group.nums[1] === ratio) {
          count += (group.counts[0] * group.counts[1])
        }
        if (num === group.nums[1]) group.counts[1]++
      }
    }
    
    if (!inGroup) {
      groups.push({ nums: [num], counts: [1, 0] })
    }
  }
  
  return count
}

Aaaaaand I made it worse. Whoops. In addition to not helping the speed my algorithm ended up failing some test cases.

The gist of what I was trying to accomplish was this:

• Store an array of all pairs + possible pairs as I iterated through the array.
• Since I had both pairs and incomplete pairs, I needed to check which one it was.
• If the number being checked would make a triplet, add to the count however many pairs had been made so far.
• I wanted to keep the amount of entries in groups to a minimum so I added a count for the first digit in the pair. This count would be incremented if the number being checked was the same.

I realized that last point is what caused a problem. It would cause the count to be incremented more than it should be. Plus, the code just looks kind of gross. So I tried to make it better.

const solve = (nums, ratio) => {
  const groups = []
  let count = 0
  
  for (const num of nums) {
    if (num % ratio !== 0 && num !== 1) continue
    
    for (const group of groups) {
      if (group.nums.length === 1) {
        if (num / group.nums[0] === ratio) {
          group.nums.push(num)
        }
      } else {
        if (num / group.nums[1] === ratio) count += group.count
        if (num === group.nums[1]) group.count++
      }
    }
    
    groups.push({ nums: [num], count: 1 })
  }
  
  return count
}

Unfortunately, I did not make it better. The code looks better but it still failed a test and still timed out. I needed to rethink how I was approaching this problem.

Iteration 3: using dictionaries but backwards!

I tried to do 2 things with my next attempt:

• Cut down the times I need to iterate through an array. My previous “dictionary” wasn’t really one, since it was an array. Somehow I overlooked that.
• Reduce the amount of times I do math on the numbers for checking conditions.

I came up with this:

const solve = (nums, ratio) => {
  const numCounts = {}
  const pairCounts = {}
  let count = 0
  
  for (let i = nums.length - 1; i >= 0; i--) {
    const [val, nextVal] = [nums[i], nums[i] * ratio]
    
    if (nextVal in pairCounts) {
      count += pairCounts[nextVal]
    }
    
    if (nextVal in numCounts) {
      pairCounts[val] = (pairCounts[val] || 0) + numCounts[nextVal]
    }
    
    numCounts[val] = (numCounts[val] || 0) + 1
  }
  
  return count
}

And hey, it worked perfectly! I realized I was really overcomplicating things before. Going backwards helped keep things clean and fast. Here’s the gist of it:

• Use 2 dictionaries. One dictionary keeps a count of numbers encountered so far and the other keeps a count of pairs that have been made with any given number.
• Every iteration, check if the current number could make a triplet out of existing pairs (this would be the smallest part of the triplet, so it’s a simple check if is in the pairCounts dictionary).
• After that, check if the current number could make a pair out of existing numbers, this is the same check as the previous step but using numCounts instead of pairCounts
• Then just add to the count

And like that, problem solved!

TL:DR – My Final Solution

const solve = (nums, ratio) => {
  const numCounts = {}
  const pairCounts = {}
  let count = 0
  
  for (let i = nums.length - 1; i >= 0; i--) {
    const [val, nextVal] = [nums[i], nums[i] * ratio]
    
    if (nextVal in pairCounts) {
      count += pairCounts[nextVal]
    }
    
    if (nextVal in numCounts) {
      pairCounts[val] = (pairCounts[val] || 0) + numCounts[nextVal]
    }
    
    numCounts[val] = (numCounts[val] || 0) + 1
  }
  
  return count
}