Read-only web browser

Solutions

Browse the interview solutions folder with syntax highlighting.

last-stone-weight.py

heap/last-stone-weight.py · Python · 2.1 KiB · 2026-02-14 11:00

Back to folder
# Brute force, max heap, bucket sort, binary search on insert.
# heap

# Brute force - O(n^2 log n) time, O(1) space
def lastStoneWeight(stones: list[int]) -> int:

    while len(stones) > 1:
        stones.sort()

        first = stones.pop()
        second = stones.pop()
        if first > second:
            stones.append(first - second)

    return stones[0] if stones else 0


# Max heap - O(n log n) time, O(1) space
def lastStoneWeight(stones: list[int]) -> int:
    import heapq

    stones = [-s for s in stones]
    heapq.heapify(stones)

    while len(stones) > 1:
        first = heapq.heappop(stones) * -1
        second = heapq.heappop(stones) * -1

        if first > second:
            heapq.heappush(stones, abs(first - second) * -1)

    return abs(stones[0]) if stones else 0


# Bucket sort - O(n + w) time, O(w) space
def lastStoneWeight(stones: list[int]) -> int:

    maxStone = max(stones)
    bucket = [0] * (maxStone + 1)
    for stone in stones:
        bucket[stone] += 1

    first = second = maxStone
    while first > 0:
        if bucket[first] % 2 == 0:
            first -= 1
            continue

        j = min(first - 1, second)
        while j > 0 and bucket[j] == 0:
            j -= 1

        if j == 0:
            return first
        second = j
        bucket[first] -= 1
        bucket[second] -= 1
        bucket[first - second] += 1
        first = max(first - second, second)
    return first


# Binary search on insert - O(n^2) time, O(1) space
def lastStoneWeight(stones: list[int]) -> int:
    stones.sort()
    n = len(stones)

    while n > 1:
        cur = stones.pop() - stones.pop()
        n -= 2
        if cur > 0:
            l, r = 0, n
            while l < r:
                mid = (l + r) // 2
                if stones[mid] < cur:
                    l = mid + 1
                else:
                    r = mid
            pos = l
            n += 1
            stones.append(0)
            for i in range(n - 1, pos, -1):
                stones[i] = stones[i - 1]
            stones[pos] = cur

    return stones[0] if n > 0 else 0