Equalizing Array Element

import sys, bisect, collections, heapq, functools
input = sys.stdin.readline


rs = lambda: input().strip()
ri = lambda: int(input())
rmi = lambda: map(int, input().split())
ra = lambda: [int(x) for x in input().split()]

INF = 10**18
MOD = 10**9 + 7

def solve():
    """
    Divide each number in the array by d until 0
    Use dict to save the number of steps to reach each smaller number
    Use heap to keep only k smallest steps
    Iterate through all smaller numbers and take global min step to reach them
    """
    freq = collections.defaultdict(list)
    for x in A:
        cnt = 0
        while x:
            heapq.heappush(freq[x], cnt)
            cnt += 1
            x //= d
    ans = INF
    for h in freq.values():
        tmp = 0
        if len(h) >= threshold:
            t = threshold
            while t:
                tmp += heapq.heappop(h)
                t -= 1
            ans = min(ans, tmp)
    return ans


test_case = 1
for _ in range(test_case):
    threshold, d = 3, 2
    A = [1,2,3,4,5]
    print(solve())