Python三路快排最后两个MLE？

是python的问题还是代码的问题...?

from random import randint

# 在前后指针法的基础上，将数列分成大于、等于、小于三个部分
def quick_sort3(arr, s, l, r):
    if l >= r:
        return
    randindex = randint(l, r)
    arr[l], arr[randindex] = arr[randindex], arr[l] # 交换枢轴值与第一个数
    i, j = l, l + 1 
    k = r + 1
    while j < k:
        if arr[j] < arr[l]:
            i += 1
            arr[i], arr[j] = arr[j], arr[i]
            j += 1
        elif arr[j] > arr[l]:
            k -= 1
            arr[j], arr[k] = arr[k], arr[j]
        else:
            j += 1
    arr[l], arr[i] = arr[i], arr[l]
    if s < i:
        quick_sort3(arr, s, l, i - 1)
    elif s >= k:
        quick_sort3(arr, s, k, r)
    else:
        return

n,s = map(int, input().split())
nums = [int(_) for _ in input().split()]
quick_sort3(nums, s, 0, n-1)
print(nums[s])