Computer >> บทช่วยสอนคอมพิวเตอร์ >  >> การเขียนโปรแกรม >> Python

โปรแกรมหาห้องที่ใกล้ที่สุดจาก Queries ใน Python

ปัญหา

กำหนดอาร์เรย์ rooms แต่ละอันคือ [roomId, size] และ queries แต่ละอันคือ [preferred, minSize] ต้องหาห้องที่ size ≥ minSize และ |roomId - preferred| น้อยที่สุด ถ้าเท่ากันเลือก roomId น้อยกว่า ไม่มีให้ -1

ตัวอย่าง

Rooms: [[2,2],[1,2],[3,2]] Queries: [[3,1],[3,3],[5,2]]

Output: [3, -1, 3]

อัลกอริทึม (Offline + Binary Search)

  1. เรียง Rooms ตาม Size แล้ว RoomId
  2. แปลง Queries เป็น (preferred, minSize, index) เรียงตาม MinSize ใหญ่ไปเล็ก
  3. ใช้ Sorted List (bisect) เก็บ RoomId ที่ผ่าน MinSize
  4. สำหรับแต่ละ Query: เพิ่ม Rooms ที่ size ≥ minSize ลง Sorted List แล้ว Binary Search หา RoomId ใกล้ที่สุด

ตัวอย่างโค้ด

import bisect

def solve(rooms, queries):
    rooms.sort(key=lambda x: (x[1], x[0]))
    queries = [(qid, size, i) for i, (qid, size) in enumerate(queries)]
    queries.sort(key=lambda x: (x[1], x[0], x[2]), reverse=True)
    
    ans = [-1] * len(queries)
    X = []  # Sorted list of room IDs
    
    for qid, size, i in queries:
        while rooms and rooms[-1][1] >= size:
            idr, _ = rooms.pop()
            bisect.insort(X, idr)
        
        if X:
            j = bisect.bisect(X, qid)
            if j == len(X):
                ans[i] = X[-1]
            elif j == 0:
                ans[i] = X[0]
            else:
                if X[j] - qid < qid - X[j-1]:
                    ans[i] = X[j]
                else:
                    ans[i] = X[j-1]
    return ans

rooms = [[2,2],[1,2],[3,2]]
queries = [[3,1],[3,3],[5,2]]
print(solve(rooms, queries))  # [3, -1, 3]

ความซับซ้อน

  • Time: O((n + q) log n)
  • Space: O(n)