786. 第k个数 - AcWing题库

Question

Content

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

输入格式

第一行包含两个整数 n 和 k。

第二行包含 n 个整数（所有整数均在 1 ~ $10^9$ 范围内），表示整数数列。

输出格式

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

数据范围

1≤n≤100000 , 1≤k≤n

输入样例：

5 3
2 4 1 5 3

输出样例：

3

Solution

Java

import java.util.Scanner;
import java.util.stream.IntStream;

public static void main(String[] args) {
    // Create a Scanner object to read input from the console
    Scanner scanner = new Scanner(System.in);

    // Read two integers, 'n' and 'k', from the user
    int n = scanner.nextInt(); // Total number of elements
    int k = scanner.nextInt(); // The k-th smallest element to find

    // Create an array 'nums' to store 'n' integers input by the user
    int[] nums = IntStream.range(0, n)
            .map(i -> scanner.nextInt()) // Read 'n' integers from the user and store them in the array
            .toArray();

    // Sort the 'nums' array using the Quick Sort algorithm
    quickSort(nums, 0, n - 1);

    // Print the k-th smallest element (0-based index) from the sorted array
    System.out.println(nums[k - 1]);
}

// Quick Sort algorithm to sort an integer array 'nums' between indices 'left' and 'right'
private static void quickSort(int[] nums, int left, int right) {
    // Base case: If the left index is greater than or equal to the right index, the array is already sorted
    if (left >= right) {
        return;
    }

    // Initialize two pointers 'i' and 'j' and choose a pivot element 'p' (median of the first and last elements)
    int i = left - 1;
    int j = right + 1;
    int p = nums[(left + right) >> 1]; // Find the pivot element

    // Partition the array 'nums' into two subarrays based on the pivot element 'p'
    while (i < j) {
        // Move the left pointer 'i' to the right until 'nums[i]' is greater than or equal to 'p'
        while (nums[++i] < p) ;

        // Move the right pointer 'j' to the left until 'nums[j]' is less than or equal to 'p'
        while (nums[--j] > p) ;

        // If 'i' is still less than 'j', swap the elements at 'i' and 'j' to maintain the partition
        if (i < j) {
            swap(nums, i, j);
        }
    }

    // Recursively sort the left and right subarrays
    quickSort(nums, left, j);
    quickSort(nums, j + 1, right);
}

// Helper function to swap two elements in an array
private static void swap(int[] nums, int left, int right) {
    // Swap elements using XOR to avoid a temporary variable
    nums[left] = nums[left] ^ nums[right];
    nums[right] = nums[left] ^ nums[right];
    nums[left] = nums[left] ^ nums[right];
}