三种语言实现快速选择（C++/Python/Java)

题目

给定一个长度为 𝑛 的整数数列，以及一个整数𝑘，请用快速选择算法求出数列从小到大排序后的第 𝑘 个数。

输入格式

第一行包含两个整数 𝑛 和 𝑘。

第二行包含 𝑛 个整数（所有整数均在 1∼109 范围内），表示整数数列。

输出格式

输出一个整数，表示数列的第 𝑘 小数。

数据范围

1≤n≤100000
1≤k≤n

输入样例：

5 3
2 4 1 5 3

输出样例：

3

C++

#include <bits/stdc++.h>

using namespace std;

const int N = 100010;
int a[N];

int quick_select(int a[], int n, int k, int l, int r)
{
    if(l >= r)
        return a[l];
    int i = l-1, j = r+1;
    int x = a[(l+r)>>1];
    while(i < j)
    {
        do i++; while(a[i] < x);
        do j--; while(a[j] > x);
        if(i < j)
            swap(a[i], a[j]);
    }
    if(j >= k)
        return quick_select(a, n, k, l, j);
    else
        return quick_select(a, n, k, j+1, r);
}

int main()
{
    int n,k;
    cin >> n >> k;
    for(int i = 1; i <= n; i++)
        cin >> a[i];
    cout << quick_select(a, n, k, 1, n);
}

Python

n, k = map(int,(input().split()))
a = [int(x) for x in input().split()]

def quick_select(a,  k, l, r):
    if l >= r:
        return 
    i, j = l-1, r+1
    x = a[(l+r)>>1]
    while i < j:
        while True:
            i += 1
            if a[i] >= x:
                break 
        while True:
            j -= 1
            if a[j] <= x:
                break 
        if i < j:
            a[i], a[j] = a[j], a[i]
    
    if j >= k:
        quick_select(a, k, l, j)
    else:
        quick_select(a, k, j+1, r)

quick_select(a, k-1, 0, n-1)
print(a[k-1])

Java

import java.util.*;

public class Main{
    public static void main(String[] args){
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        int k = sc.nextInt();
        int[] a = new int[n];
        for(int i = 0; i < n; i++)
            a[i] = sc.nextInt();
        System.out.print(quick_select(a, 0, n-1, n, k-1));
    }
    public static int quick_select(int[] a, int l, int r, int n, int k){
        if(l >= r)
            return a[l];
        int i = l-1, j = r+1;
        int x = a[(l+r) >>1];
        while(i < j)
        {
            while(a[++i] < x);
            while(a[--j] > x);
            if(i < j)
            {
                int t = a[i];
                a[i] = a[j];
                a[j] = t;
            }
        }
        if(j >= k)
            return quick_select(a, l, j, n, k);
        else
            return quick_select(a, j+1, r, n, k);
    }
}