从有序序列中求最大 不大于目标 的数的下标的二分查找怎么写比较优雅？

Question

本身这个题挺简单的，但是如果增加这样一个要求怎么写：
如果序列中有多个等于目标的数，则可以传入一个flag参数，来决定返回等于目标的数最大下标还是最小下标

举个例子：
binsearch([1, 3, 3], 3, "max") 返回 2
binsearch([1, 3, 3], 3, "min") 返回 1
binsearch([1, 1], 2, "max") 和 binsearch([1, 1], 2, "min") 都返回 1
就是只有当数列中有多个等于目标的数时flag参数才有意义。

我写的只考虑flag=="min"的代码，我觉得最后几行if已经很不优雅了，flag=="max"怎么也写不对

def binsearch(a, target, flag="min"):
    ''' Find the index of the max one not greater than target'''
    l, r = 0, len(a)-1
    while (l < r-1):
        mid = l + (r-l) / 2
        if a[mid] >= target:
            r = mid
        else:
            l = mid+1
    if a[l] == target: return l
    if a[r] == target: return r
    if a[r] < target: return r
    if a[l] < target: return l
    return -1

assert(binsearch([1, 3], 1) == 0)
assert(binsearch([1, 3], 3) == 1)
assert(binsearch([1, 3], 2) == 0)
assert(binsearch([1, 1], 2) == 1)
assert(binsearch([1, 1, 3], 1) == 0)
assert(binsearch([2, 2], 1) == -1)
assert(binsearch([1, 3, 5], 1) == 0)
assert(binsearch([1, 3, 5], 5) == 2)
assert(binsearch([1, 3, 5, 5], 5) == 2)
assert(binsearch([1, 1, 1, 2, 2, 2, 3, 4, 5, 5], 2) == 3)


# assert(binsearch([1, 3], 1, "max") == 0)
# assert(binsearch([1, 3], 3, "max") == 1)
# assert(binsearch([1, 3], 2, "max") == 0)
# assert(binsearch([1, 1], 2, "max") == 0)
# assert(binsearch([1, 1, 3], 1, "max") == 1)
# assert(binsearch([2, 2], 1, "max") == -1)
# assert(binsearch([1, 3, 5], 1, "max") == 0)
# assert(binsearch([1, 3, 5], 5, "max") == 2)
# print binsearch([1, 3, 5, 5], 5, "max")
# assert(binsearch([1, 3, 5, 5], 5, "max") == 3)

Answer 1

我后来又想了想，同时参考了楼上@brayden 的思想，把两种功能的实现完全分开，重新写了一个

def binsearchmax(a, target):
    l, r = 0, len(a)-1
    while (l < r-1):
        mid = l + (r-l) / 2
        if a[mid] <= target:
            l = mid
        else:
            r = mid-1
    if a[r] <= target: return r
    if a[l] <= target: return l
    return -1

def binsearchmin(a, target):
    l, r = 0, len(a)-1
    while (l < r-1):
        mid = l + (r-l) / 2
        if a[mid] < target:
            l = mid
        else:
            r = mid
    if a[l] == target: return l
    if a[r] == target: return r
    if a[r] < target: return r
    if a[l] < target: return l
    return -1

def binsearch(a, target, flag="min"):
    if flag == "min": 
        return binsearchmin(a, target)
    else: 
        return binsearchmax(a, target)

assert(binsearch([1, 3], 1) == 0)
assert(binsearch([1, 3], 3) == 1)
assert(binsearch([1, 3], 2) == 0)
assert(binsearch([1, 1], 2) == 1)
assert(binsearch([1, 1, 3], 1) == 0)
assert(binsearch([2, 2], 1) == -1)
assert(binsearch([1, 3, 5], 1) == 0)
assert(binsearch([1, 3, 5], 5) == 2)
assert(binsearch([1, 3, 5, 5], 5) == 2)
assert(binsearch([1, 1, 1, 2, 2, 2, 3, 4, 5, 5], 2) == 3)


assert(binsearch([1, 3], 1, "max") == 0)
assert(binsearch([1, 3], 3, "max") == 1)
assert(binsearch([1, 3], 2, "max") == 0)
assert(binsearch([1, 1], 2, "max") == 1)
assert(binsearch([1, 1, 3], 1, "max") == 1)
assert(binsearch([2, 2], 1, "max") == -1)
assert(binsearch([1, 3, 5], 1, "max") == 0)
assert(binsearch([1, 3, 5], 5, "max") == 2)
assert(binsearch([1, 3, 5, 5], 5, "max") == 3)

我自己构造的测试数据全部通过了。总结一下，我觉得我一开始试图把两种功能的逻辑混到一起的这个思想是导致我写不出来的一个重要原因。再次感谢楼上的@brayden

Answer 2

嗯要是偷点儿懒的话开始之前往 a 的最前面插 -无穷，最后面插 +无穷，不过这个要求确实很难精简最后四句 if。

update. 想了想可以写一起。。。

def binsearch(a, target, flag="min"):
    if a[0] > target:
        return -1
    a = [-999099] + a + [999099]
    ismin = (flag == "min")
    l, r = 0, len(a) - 1
    while (l < r - 1):
        mid = (l + r) / 2
        if a[mid] > target or ismin and a[mid] == target:
            r = mid
        else:
            l = mid
    if ismin:
        return r - 1 if a[r] == target else l - 1
    else:
        return l - 1

Answer 3

把两种需求写在一个函数里, 逻辑不清楚也给自己找麻烦. 我写的话, 两个函数分别计算max, min. 根据 输入决定去调用哪一个. java代码:

private int findLower(int[] A, int target) {
    int len = A.length;
    int low = 0, high = len - 1;
    while(low <= high){
        int mid = ((high - low) >> 1) + low;
        if((A[mid] < target) && (mid == len - 1 || A[mid + 1] >= target))
            return mid;
        if(A[mid] < target)
            low = mid + 1;
        else // A[mid] >= target
            high = mid - 1;
    }
    return -1;
}

private int findHigher(int[] A, int target) {
    int len = A.length;
    int low = 0, high = len - 1;
    while(low <= high){
        int mid = ((high - low) >> 1) + low;
        if((mid == 0 || A[mid - 1] <= target) && (A[mid] > target))
            return mid;
        if(A[mid] <= target)
            low = mid + 1;
        else // A[mid] > target
            high = mid - 1;
    }
    return len;
}

public int binSearch(int[] A, int target, boolean lower){
    return lower ? findLower(A, target) : findHigher(A, target);
}