Enable Javascript in your browser and then refresh this page, for a much enhanced experience.
101-liner: 10**12 in 2min solution in Speedy category for Ugly Numbers by przemyslaw.daniel
from heapq import heappush, heappushpop, nlargest
from itertools import count
from math import log
class MaxHeap:
""" Keep only 'size' largest values """
def __init__(self, size: int):
self.heap = []
self.size = size
def add(self, element):
if len(self.heap) < self.size:
heappush(self.heap, element)
else:
heappushpop(self.heap, element)
def get(self):
return nlargest(self.size, self.heap)[::-1]
def ugly_number_small(position: int) -> int:
""" Regular calculations for small cases """
found, ugly = {1}, 1
for _ in range(position):
found -= {ugly := min(found)}
found |= {2*ugly, 3*ugly, 5*ugly}
return ugly
def ugly_final(value: int, number: int) -> int:
""" At this point we have a value which is close to the nth ugly number
Now we need to count all ugly numbers below the value, collect top once
and after sorting read the final value by calculating the proper offset.
Due to floating point operations this function might be eventually
not precise enough and can be replaced by for example decimal module or
integer calculations """
size2 = round(log(value, 2)) + 1
size3 = round(log(value, 3)) + 1
size5 = round(log(value, 5)) + 1
found = MaxHeap(size5)
step2, step3, step5 = log(2), log(3), log(5)
result, value, base2 = 0, log(value), 0
for i in range(size2):
base23, best = base2, (0, (0, 0, 0))
for j in range(size3):
if value >= base23:
k = int((value - base23) / step5)
found.add((base23 + k*step5, (i, j, k)))
result += k + 1
else:
break
base23 += step3
base2 += step2
_, (a, b, c) = found.get()[number-result-1]
return 2**a*3**b*5**c
def ugly_estimate(value: int) -> int:
""" Estimate amount of ugly numbers below given value. It is assured
that ugly_number(ugly_estimate(value))) <= value """
x, y, z = log(value, 2), log(value, 3), log(value, 5)
a = (x + 1) * (y + 1) * (z + 1) / 6
b = (x + 2) * (y + 2) * (z + 2) / 6
return int((a + b - z - 1) / 2)
def find_closest_to_ugly(number: int) -> int:
""" find the closest value of the nth ugly number using binary search """
lower = 1
upper = next(c for x in count() if ugly_estimate(c := 2**x) > number)
while True:
middle = (lower + upper) // 2
ue = ugly_estimate(middle)
if ue == number:
return middle
elif ue < number:
lower = middle
else:
upper = middle
def ugly_number(number: int) -> int:
""" This is a simple method that performs the calculations in
two steps. Firstly, we look for a value that is as close
as possible to the nth ugly number. Secondly, we find
the correct solution """
if number < 99:
# handle small cases where calculations are not proper
return ugly_number_small(number)
value = find_closest_to_ugly(number)
return ugly_final(value, number)
Jan. 26, 2022
Comments:
