ข้อ 50 · LC2336 Smallest Number in Infinite Set (เลขน้อยสุดเซ็ตอนันต์) 🟡

โจทย์ (LC2336): ให้ออกแบบ class ชื่อ SmallestInfiniteSet ที่แทนเซ็ตซึ่งบรรจุ positive integer (จำนวนเต็มบวก) ทุกตัวตั้งแต่ 1, 2, 3, ... ไปจนถึงอนันต์ตั้งแต่เริ่มต้น โดยต้อง support (รองรับ) สอง method คือ popSmallest() ซึ่ง remove และ return ค่าที่น้อยที่สุดที่ยังอยู่ในเซ็ต และ addBack(num) ซึ่งเพิ่มจำนวนเต็มบวก num กลับเข้าเซ็ต ถ้ามันยังไม่อยู่ในเซ็ตอยู่แล้ว

แนวทาง — ต้องใช้อะไร & คิดยังไง

โครงสร้างที่ใช้: counter (ตัวนับ) current หนึ่งตัว + min-heap หนึ่งอัน + set กันเลขซ้ำ เราไม่ต้อง store เลข infinity จริง ๆ เพราะเลขในช่วง current, current+1, ... ยัง sorted (เรียงเป็นระเบียบ) อยู่แล้ว แค่ track (จำ) ว่าถึงไหนก็พอ

คิดแบบง่าย/ช้าก่อน: ถ้า store เลขทั้งหมดจริง ๆ จะเป็นไปไม่ได้เพราะ infinity ปัญหาเดียวที่ต้อง handle แยกคือเลขที่ถูก addBack กลับมา ซึ่งอาจเล็กกว่า current เราจึง store มันใน min-heap เพื่อให้ pop ตัว minimum ออกก่อนได้เสมอ

1. initialize current = 1 (เลขต่อไปในสาย infinity ที่ยังไม่เคย pop), added = min-heap ว่าง, in_heap = set ว่าง
2. popSmallest: ถ้า heap มีของ (การันตีว่าเล็กกว่า current) pop จาก heap ก่อน แล้ว remove ออกจาก set
3. ถ้า heap ว่าง pop current แล้ว increment current ขึ้นหนึ่ง
4. addBack(num): เพิ่มกลับได้เฉพาะเลขที่ถูก pop ไปแล้ว (num < current) และยังไม่อยู่ใน heap จึง push เข้า heap และ add เข้า set

ไล่ทีละสเต็ป

จำลอง sequence ของ operation pop, pop, pop, addBack(2), pop, pop, pop:

import heapq

class SmallestInfiniteSet:
    def __init__(self):
        self.current = 1      # เลขต่อไปในสาย 1,2,3,... ที่ยังไม่เคยหยิบ
        self.added = []       # min-heap ของเลขที่ถูก addBack กลับมา
        self.in_heap = set()  # กันไม่ให้ heap มีเลขซ้ำ

    def popSmallest(self):
        # ถ้ามีเลขที่ addBack กลับมา และมันเล็กกว่า current ให้หยิบจาก heap ก่อน
        if self.added:
            x = heapq.heappop(self.added)
            self.in_heap.discard(x)
            return x
        # ไม่งั้นหยิบตัวถัดไปจากสายอนันต์
        x = self.current
        self.current += 1
        return x

    def addBack(self, num):
        # เพิ่มกลับได้เฉพาะเลขที่ถูกหยิบออกไปแล้ว (num < current) และยังไม่อยู่ใน heap
        if num < self.current and num not in self.in_heap:
            heapq.heappush(self.added, num)
            self.in_heap.add(num)

s = SmallestInfiniteSet()
print(s.popSmallest())  # 1
print(s.popSmallest())  # 2
print(s.popSmallest())  # 3
s.addBack(2)
print(s.popSmallest())  # 2
print(s.popSmallest())  # 4
print(s.popSmallest())  # 5

1
2
3
2
4
5