Count Triplets

  • + 0 comments

    This solution is extremely clever — or at least, it took me a while to understand it. :)

    In case it is of use to anyone else, I rewrote the code with inline comments and descriptive variable names, which helped me to get it. I refer to a harmonic pair as two elements where the leftmost element is the base and the rightmost element is the head, and the the head = r * base.

    def countTriplets(arr, r):
        # running count of triplets
        triplet_count = 0
        # maps element val to count of elements with that value, to our right
        val_to_count = {}
        # maps element val to count of harmonic pairs with that element 
        # as the base, where the head of the pair is to our right
        baseval_to_paircount = {}
    
        # iterate in reverse, so that when we check if an element
        # is the base of a harmonic pair, we can be sure we have already seen
        # any potential head of the pair, since it must be to our right,
        # since we only consider triplets where i < j < k so we only consider
        # pairs where i < j
        for base in reversed(arr):
            head = base * r
            # 1. determine if this pair's head is the base of a pair to the right
            #
            # In other words, check if this is the base of a triplet
            #
            # if this element's head is another pair's base ...
            if head in baseval_to_paircount:
                # ... then we found a triplet.
                # increase the triplet triplet_count, by the number of
                # pairs usig our head as their base
                triplet_count += baseval_to_paircount[head]
    
            # 2. update triplet_count of pairs with this base value
            #
            # (We must do this second, so as to not over-increment
            # triplet_coont, by adding ourselves to baseval_to_paircount
            # when base=head.)
            #
            # if this element's head has appeared to the right in the array ...
            if head in val_to_count:
                # ... then increase the triplet_count of pairs starting at our current base
                # by the number of times head has appeared so far
                baseval_to_paircount[base] = baseval_to_paircount.get(base, 0) + val_to_count[head]
    
            # 3. Note this element value appeared, for future iterations
            # examining elements to the left
            #
            # (We must do this last, so we don't over-increment
            # baseval_to_paircount by coounting ourselves as our own head
            # when base=head.)
            #
            val_to_count[base] = val_to_count.get(base, 0) + 1
        return triplet_count