劍指offer：在排序數組中查找數字

題目:
統計一個數字在排序數組中出現的次數. 例如輸入排序數組{1,2,3,3,3,3,4,5},因爲3在這個數中出現了4次,輸出4.

# -*- coding: utf-8 -*-
# @Time         : 2019-07-13 15:10
# @Author       : Jayce Wong
# @ProjectName  : job
# @FileName     : getNumberOfK.py
# @Blog         : https://blog.51cto.com/jayce1111
# @Github       : https://github.com/SysuJayce

def getFirstK(data, k):
    start, end = 0, len(data) - 1
    while start <= end:
        mid = (start + end) >> 1
        if data[mid] == k:
            # 關鍵在於，若是mid是k，那麼就判斷前一個元素是否是也是k，若是是，說明這個位置不是
            # 第一次出現，要在左邊繼續查找。不然，直接返回mid，由於是第一次出現了
            if mid - 1 >= start and data[mid - 1] == k:
                end = mid - 1
            else:
                return mid
        elif data[mid] < k:
            start = mid + 1
        else:
            end = mid - 1

    return -1

def getLastK(data, k):
    start, end = 0, len(data) - 1
    while start <= end:
        mid = (start + end) >> 1
        if data[mid] == k:
            if mid + 1 <= end and data[mid + 1] == k:
                start = mid + 1
            else:
                return mid
        elif data[mid] < k:
            start = mid + 1
        else:
            end = mid - 1

    return -1

def getNumberOfK(data, k):
    """
    要獲取一個有序數組中某個元素出現的次數，最直觀的作法就是遍歷整個數組，而後統計該元素的出現次數，
    這樣作的時間複雜度是O(n)

    可是因爲這個數組是有序的，咱們能夠考慮利用二分查找的方法來解決這個問題。
    若是咱們先利用二分查找定位到了這個元素，而後再往前日後遍歷，這樣的話時間複雜度也仍是O(n)。

    可是若是咱們在利用二分查找的時候，想辦法定位這個元素第一次出現的下標和最後一次出現的下標。
    在利用二分查找找到一個這個元素以後，判斷這個元素是不是第一個，也就是對比這個元素的前一個是否也
    是k，若是不是，說明這個元素就是第一個元素，不然在這個下標的左邊繼續查找。
    對於最後一次出現的下標同理。
    """
    if not data:
        return 0
    first = getFirstK(data, k)
    last = getLastK(data, k)
    if first != -1 and last != -1:
        return last - first + 1
    else:
        return 0

def main():
    data = [1, 2, 3, 3, 3, 3, 4, 5]
    k = 3
    print(getNumberOfK(data, k))

if __name__ == '__main__':
    main()