kth-two-sorted-arrays-invariant

Answer

The invariant. A cut splits both arrays: take i elements from A and j = k − i from B into the left half, so leftA + leftB == k. The cut is a valid k-th partition iff every left element is ≤ every right element:

max(A[i-1], B[j-1])  <=  min(A[i], B[j])

Since each array is individually sorted, this reduces to exactly two boundary checks, and both must be non-strict (<=):

A[i-1] <= B[j]   and   B[j-1] <= A[i]        (with ±inf sentinels at the edges)

The bug. A mixed <=/< predicate, e.g. A[i-1] < B[j] and B[j-1] <= A[i], rejects a cut when a duplicate value straddles the boundary (A[i-1] == B[j]). But that cut is genuinely valid — the k-th element can legitimately sit on either side of the boundary. The search then pushes i the wrong way. Because j = k − i must stay in [0, n], the feasible window i ∈ [max(0, k−n), min(m, k)] shrinks to a handful of values when k is near an array end; the binary search exhausts the window without ever accepting the true cut, and falls back on a partition whose left maximum exceeds the right minimum, returning an element that is provably not the k-th smallest.

The fix. Test both sides with <=, and clamp the binary-search window so the invariant leftA + leftB == k is satisfiable for every k ∈ [1, m+n]:

import math

def kth_smallest(a, b, k):
    """k-th smallest (1-indexed) of two sorted arrays, O(log(min(m, n)))."""
    m, n = len(a), len(b)
    if m > n:                      # binary-search the shorter array
        a, b = b, a
        m, n = n, m

    lo, hi = max(0, k - n), min(m, k)      # keeps j = k - i in [0, n]

    while lo <= hi:
        i = (lo + hi) // 2
        j = k - i

        a_left  = a[i - 1] if i > 0 else -math.inf
        a_right = a[i]     if i < m else  math.inf
        b_left  = b[j - 1] if j > 0 else -math.inf
        b_right = b[j]     if j < n else  math.inf

        # FIXED boundary predicate: max(left) <= min(right), both <=
        if a_left <= b_right and b_left <= a_right:
            return max(a_left, b_left)          # the k-th smallest
        if a_left > b_right:
            hi = i - 1                          # take fewer from A
        else:                                   # b_left > a_right
            lo = i + 1                          # take more from A

    raise ValueError("unreachable for 1 <= k <= m+n with sorted input")

Why this is now correct: max(left) <= min(right) is equivalent to a_left <= b_right and b_left <= a_right (the other two cross terms are guaranteed by per-array sortedness), a valid cut exists for every k ∈ [1, m+n], and the movement rules are monotone-safe — if a_left > b_right no larger i can be valid; if b_left > a_right no smaller i can be valid. So the search provably terminates on a valid cut and returns max(left), which is exactly the k-th smallest.