IT博客汇
  • 首页
  • 精华
  • 技术
  • 设计
  • 资讯
  • 扯淡
  • 权利声明
  • 登录 注册

    Effectively fetch the smallest element in a heap of Python

    RobinDong发表于 2023-02-24 09:56:10
    love 0

    To solve Leetcode #1675, I wrote the below code with the help of hints:

    import sys
    import heapq
    
    class Solution:
        def minimumDeviation(self, nums: List[int]) -> int:
            n = len(nums)
            heapq.heapify(nums)
            _max = max(nums)
            ans = max(nums) - min(nums)
            while True:
                item = heapq.heappop(nums)
                if item % 2 == 1:
                    heapq.heappush(nums, item * 2)
                    _max = max(_max, item * 2)
                    ans = min(ans, _max - heapq.nsmallest(1, nums)[0])
                else:
                    heapq.heappush(nums, item)
                    break
            print("stage1:", nums)
            nums = [-item for item in nums]
            heapq.heapify(nums)
            _max = max(nums)
            while True:
                item = heapq.heappop(nums)
                if item % 2 == 0:
                    heapq.heappush(nums, item // 2)
                    _max = max(_max, item // 2)
                    ans = min(ans, _max - heapq.nsmallest(1, nums)[0])
                else:
                    break
                    
            return ans      

    I know, the code looks quite messy. But the more annoying problem is: it exceeded the time limit.

    After learning the solutions from these smart guys, I realized how stupid I am — I can just use nums[0] instead of heapq.nsmallest(1, nums)[0] to get the smallest element in a heap, just as the official document said.

    Then I just change two lines of my code and it passed all the test cases in time:

    import sys
    import heapq
    
    class Solution:
        def minimumDeviation(self, nums: List[int]) -> int:
            n = len(nums)
            heapq.heapify(nums)
            _max = max(nums)
            ans = max(nums) - min(nums)
            while True:
                item = heapq.heappop(nums)
                if item % 2 == 1:
                    heapq.heappush(nums, item * 2)
                    _max = max(_max, item * 2)
                    ans = min(ans, _max - nums[0])
                else:
                    heapq.heappush(nums, item)
                    break
            print("stage1:", nums)
            nums = [-item for item in nums]
            heapq.heapify(nums)
            _max = max(nums)
            while True:
                item = heapq.heappop(nums)
                if item % 2 == 0:
                    heapq.heappush(nums, item // 2)
                    _max = max(_max, item // 2)
                    ans = min(ans, _max - nums[0])
                else:
                    break
                    
            return ans


沪ICP备19023445号-2号
友情链接