Tìm kiếm nhị phân: loại một nửa sau mỗi lần hỏi

Lan viết công cụ tra điểm cho kỳ thi thử có 100 000 thí sinh. Mỗi thí sinh nhập điểm của mình và hỏi: "Có bao nhiêu bạn điểm thấp hơn mình?". Chương trình của Lan duyệt cả danh sách cho mỗi câu hỏi. Khi 100 000 thí sinh cùng hỏi, máy phải làm khoảng 10¹⁰ phép so sánh và bị Quá thời gian. Danh sách điểm đã được sắp xếp sẵn, nhưng chương trình của Lan chưa tận dụng điều đó. Bài này dạy cách trả lời mỗi câu hỏi chỉ với khoảng 17 lần so sánh.

Mỗi lần so sánh loại được một nửa

Hãy nghĩ tới trò đoán số từ 1 đến 100. Người đoán giỏi luôn hỏi số ở giữa. Nếu câu trả lời là "lớn hơn", cả nửa dưới bị loại; nếu là "nhỏ hơn", cả nửa trên bị loại. Sau mỗi câu hỏi, số khả năng còn lại giảm một nửa.

Đoạn chương trình dưới đây đếm số lần phải chia đôi cho đến khi chỉ còn một khả năng:

for n in [10, 1000, 100_000, 1_000_000_000]:
    steps = 0
    size = n
    while size > 1:
        size = (size + 1) // 2
        steps += 1
    print(n, steps)

10 4
1000 10
100000 17
1000000000 30

Với một tỉ phần tử, chỉ cần 30 lần chia đôi. Số lần này xấp xỉ log₂ n, nên mỗi lần tìm kiếm nhị phân tốn O(log n) bước, so với O(n) bước khi duyệt từng phần tử. Với bài toán của Lan, 100 000 câu hỏi nhân 17 bước chỉ còn khoảng 1,7 triệu phép so sánh.

Điều kiện để dùng được: dãy Sai rồi mới tới Đúng

Tìm kiếm nhị phân không chỉ dành cho "tìm một số trong mảng đã sắp xếp". Nó dùng được mỗi khi bạn có một câu hỏi Đúng/Sai, gọi là điều kiện kiểm tra (predicate). Khi đi từ trái sang phải, câu hỏi đó cho một đoạn Sai liền nhau rồi tới một đoạn Đúng liền nhau. Tính chất này gọi là tính đơn điệu (monotonicity).

Lấy dãy đã sắp xếp a và điều kiện a[i] >= 7:

Vì dãy tăng, một khi a[i] >= 7 thì mọi phần tử bên phải cũng >= 7. Hỏi ở giữa và thấy Đúng, ta biết điểm chuyển nằm ở đó hoặc bên trái. Thấy Sai, ta biết điểm chuyển nằm bên phải. Mục tiêu của tìm kiếm nhị phân là tìm vị trí Đúng đầu tiên, tức điểm chuyển từ Sai sang Đúng.

Tìm vị trí Đúng đầu tiên

Ta giữ một khoảng [lo, hi] chứa những vị trí chưa biết Đúng hay Sai, và biến ans lưu vị trí Đúng nhỏ nhất đã thấy. Ban đầu chưa thấy vị trí Đúng nào, nên ans = n, nghĩa là "không có vị trí nào thỏa".

Mỗi vòng lặp hỏi ở giữa mid = (lo + hi) // 2:

• Nếu mid Đúng: ghi ans = mid, rồi tìm tiếp bên trái với hi = mid - 1, vì có thể còn vị trí Đúng nhỏ hơn.
• Nếu mid Sai: mọi vị trí từ lo tới mid đều Sai, nên đặt lo = mid + 1.

Vòng lặp dừng khi lo > hi, tức không còn vị trí nào chưa biết. Hình dưới theo dõi việc tìm vị trí đầu tiên có a[i] >= 7. Ô xanh là vùng đã biết Đúng, ô xám là vùng đã biết Sai:

Ở bước 1, a[3] = 7 thỏa nên ans = 3 và nửa phải bị loại. Ở bước 2 và 3, a[1] = 3 và a[2] = 4 đều không thỏa nên lo tăng dần lên 3. Khi lo = 3 > hi = 2, vòng lặp dừng với ans = 3.

Viết thành hàm nhận điều kiện kiểm tra ok, ta dùng lại được cho mọi bài:

def first_true(lo, hi, ok):
    ans = hi + 1
    while lo <= hi:
        mid = (lo + hi) // 2
        if ok(mid):
            ans = mid
            hi = mid - 1
        else:
            lo = mid + 1
    return ans

a = [1, 3, 4, 7, 7, 9, 12, 15]
n = len(a)
print(first_true(0, n - 1, lambda i: a[i] >= 7))
print(first_true(0, n - 1, lambda i: a[i] >= 10))
print(first_true(0, n - 1, lambda i: a[i] >= 100))

3
6
8

Kết quả thứ ba là 8, bằng n: không có phần tử nào >= 100, nên hàm trả về vị trí ngay sau cuối dãy. Chương trình gọi hàm cần kiểm tra trường hợp này trước khi dùng a[ans].

Đếm bằng bisect_left và bisect_right

Python có sẵn module bisect cho hai câu hỏi hay gặp nhất trên dãy đã sắp xếp:

• bisect_left(a, x): vị trí đầu tiên có a[i] >= x, cũng chính là số phần tử nhỏ hơn x.
• bisect_right(a, x): vị trí đầu tiên có a[i] > x, cũng chính là số phần tử nhỏ hơn hoặc bằng x.

Hiệu bisect_right(a, x) - bisect_left(a, x) là số lần x xuất hiện. Thay hai giá trị khác nhau, ta đếm được số phần tử trong đoạn giá trị [L, R]:

from bisect import bisect_left, bisect_right

a = [1, 2, 4, 4, 4, 7, 9]
print(bisect_left(a, 4), bisect_right(a, 4))
print("số lần xuất hiện của 4:", bisect_right(a, 4) - bisect_left(a, 4))
print("số phần tử trong [2, 7]:", bisect_right(a, 7) - bisect_left(a, 2))
print("số phần tử trong [5, 6]:", bisect_right(a, 6) - bisect_left(a, 5))

2 5
số lần xuất hiện của 4: 3
số phần tử trong [2, 7]: 5
số phần tử trong [5, 6]: 0

Bài toán của Lan giờ chỉ cần sắp xếp một lần rồi gọi bisect_left cho mỗi câu hỏi:

6 3
7 2 9 7 5 10
7
1
11

import sys
from bisect import bisect_left

data = sys.stdin.read().split()
n, q = int(data[0]), int(data[1])
scores = sorted(map(int, data[2:2 + n]))
out = []
for x in map(int, data[2 + n:2 + n + q]):
    out.append(bisect_left(scores, x))
print("\n".join(map(str, out)))

2
0
6

Tổng chi phí là O(n log n) cho việc sắp xếp và O(q log n) cho các câu hỏi, thay vì O(n·q).

from bisect import bisect_left

a = [5, 1, 4, 2]
i = bisect_left(a, 5)
print(i, i < len(a) and a[i] == 5)

4 False

Số 5 nằm ở vị trí 0, nhưng bisect_left trả về 4 và kết luận không tìm thấy.

Tìm vị trí Đúng cuối cùng

Có bài cần giá trị lớn nhất còn thỏa điều kiện, với dãy kết quả là Đúng rồi mới tới Sai. Ví dụ căn bậc hai nguyên của n là số x lớn nhất có x * x <= n. Mẫu đối xứng với first_true: khi mid Đúng thì ghi lại và tìm tiếp bên phải.

def last_true(lo, hi, ok):
    ans = lo - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        if ok(mid):
            ans = mid
            lo = mid + 1
        else:
            hi = mid - 1
    return ans

for n in [0, 1, 30, 36, 999_999_999_999_999_999]:
    print(n, last_true(0, n, lambda x: x * x <= n))

0 0
1 1
30 5
36 6
999999999999999999 999999999

Vì sao không dùng int(n ** 0.5)? Số thực của Python chỉ chính xác khoảng 15–16 chữ số. Với n lớn, phép làm tròn cho kết quả lệch:

import math
n = 999_999_999_999_999_999
print(int(math.sqrt(n)))
print(math.isqrt(n))

1000000000
999999999

Đáp án đúng là 999 999 999, vì 1 000 000 000² = 10¹⁸ lớn hơn n. Python có sẵn math.isqrt cho đúng việc này, nhưng cách tìm vị trí Đúng cuối cùng áp dụng cho mọi điều kiện, không chỉ căn bậc hai.

Tìm kiếm nhị phân trên đáp án

Nhiều bài không cho mảng nào để tìm. Đề hỏi một giá trị nhỏ nhất (hoặc lớn nhất) mà ta không biết công thức tính trực tiếp, nhưng kiểm tra một giá trị cho trước thì dễ. Khi đó ta tìm kiếm nhị phân trên chính đáp án.

Ví dụ: một tiệm có ba máy in, máy thứ i in một bản mất t[i] giây, với t = [2, 3, 7]. Các máy chạy song song. Cần ít nhất bao nhiêu giây để in đủ 10 bản?

Câu hỏi "trong T giây in được đủ 10 bản không?" dễ trả lời: máy i in được T // t[i] bản, cộng lại rồi so với 10. Điều kiện này đơn điệu: nếu T giây đủ thì T + 1 giây cũng đủ.

Đáp án là 12 giây, vị trí Đúng đầu tiên. Để tìm nó, ta cần một khoảng chắc chắn chứa đáp án. lo = 1 là thời gian nhỏ nhất có thể. Nếu chỉ dùng máy nhanh nhất, sau min(t) * k giây đã in đủ k bản, nên hi = min(t) * k chắc chắn Đúng.

3 10
2 3 7

def first_true(lo, hi, ok):
    ans = hi + 1
    while lo <= hi:
        mid = (lo + hi) // 2
        if ok(mid):
            ans = mid
            hi = mid - 1
        else:
            lo = mid + 1
    return ans

n, k = map(int, input().split())
t = list(map(int, input().split()))

def enough(T):
    return sum(T // x for x in t) >= k

print(first_true(1, min(t) * k, enough))

12

Mỗi lần gọi enough tốn O(n), và số lần gọi là log₂(min(t)·k). Với n = 10⁵ máy và k = 10⁹ bản, hi khoảng 10¹⁸, nên chỉ cần khoảng 60 lần gọi, tức 6·10⁶ phép tính.

Khi gặp một bài hỏi "nhỏ nhất là bao nhiêu để làm được…" hoặc "lớn nhất là bao nhiêu mà vẫn…", hãy thử ba câu hỏi:

1. Nếu cho trước một đáp án X, có kiểm tra được nhanh là X khả thi không?
2. Nếu X khả thi, mọi giá trị lớn hơn (hoặc nhỏ hơn) có chắc khả thi không?
3. Khoảng [lo, hi] nào chắc chắn chứa đáp án?

Trả lời được cả ba câu là bạn đã có lời giải.

Lỗi thường gặp

Tóm tắt

• Tìm kiếm nhị phân cần một điều kiện đơn điệu: một đoạn Sai rồi tới một đoạn Đúng. Mỗi lần hỏi ở giữa loại được một nửa, nên tốn O(log n) lần hỏi.
• Mẫu first_true: giữ [lo, hi] là vùng chưa biết, khi Đúng thì ans = mid; hi = mid - 1, khi Sai thì lo = mid + 1. Mẫu last_true đối xứng.
• Trên dãy đã sắp xếp, bisect_left(a, x) đếm số phần tử < x, bisect_right(a, x) đếm số phần tử <= x. Hiệu của chúng đếm phần tử trong một khoảng giá trị.
• Khi đề hỏi giá trị nhỏ nhất hoặc lớn nhất và việc kiểm tra một giá trị cho trước là dễ, hãy tìm kiếm nhị phân trên đáp án.

Kiểm tra nhanh

Đọc thêm

• Tài liệu Python về module bisect, gồm cả tham số key từ Python 3.10.
• Tài liệu Python về math.isqrt, căn bậc hai nguyên chính xác.