The standard library (ie. no special imports) comes with a heapq
module that makes the algorithm O(3n + k*(2 lg n)) => O(lg n):
import math
import heapq
def minSum(num, k):
heap = [-n for n in num] # negate values for max-heap
heapq.heapify(heap)
for i in range(k):
# Find max value
max_value = heapq.heappop(heap)
# Change max value to rounded half
# use floor since we've negated the values
heapq.heappush(heap, math.floor(max_value/2))
# Calculate minimum sum
return -sum(heap) # reverse polarity again
update 1: (from comment by @raury-daulton) combine pop/push, O(3n + k*(lg n)) => O(lg n):
def minSum(num, k):
heap = [-n for n in num] # negate values for max-heap
heapq.heapify(heap)
for i in range(k):
max_value = heap[0]
heapq.heapreplace(heap, math.floor(max_value/2))
# Calculate minimum sum
return -sum(heap) # reverse polarity again
update 2: use a max heap directly O(2n + k*(lg n)) => O(lg n)
def heapreplace_max(heap, item):
"We need to implement this ourselves from primitives."
returnitem = heap[0]
heap[0] = item
heapq._siftup_max(heap, 0)
return returnitem
def minSum(num, k):
heap = num # alias for consistency with other samples
heapq._heapify_max(heap) # make it a max heap (no negation hack)
for i in range(k):
max_value = heap[0]
heapreplace_max(heap, math.ceil(max_value/2))
return sum(heap)
update 3: final optimizations (this requires the input array to be an array of int
s).
def minSum(num, k):
heap = num # alias for consistency with other samples
heapq._heapify_max(heap)
for i in range(k):
max_value = heap[0]
if max_value == 1: # exit early if there is no more work
break
new_val = (max_value >> 1) + (max_value & 1) # same as, but much faster than math.ceil(max_value/2)
heapreplace_max(heap, new_val)
return sum(heap)
The upper bound on k is sum(math.floor(math.log(v, 2)+1) for v in num)
(i.e. the total number of bits required to represent all the input numbers). It may be faster to precompute the loop variable than to have if max_value == 1:
in the loop, i.e.:
for i in range(min(k, int(sum(floor(log(v, 2)+1) for v in num)))):
max_value = heap[0]
new_val = ...
but I haven't actually measured this.