Leetcode 398 - Random Pick Index (Variation)

OG: https://leetcode.com/problems/random-pick-index/description/

Variation: select K random elements from an array

1. Modified Fisher-Yates (Knuth) shuffle - for when array can be modified

2. Reservoir sampling - for streaming data or when array modification is not desired

import java.util.Random;
import java.util.Arrays;

/**
 * Implementation of algorithms to select k random elements from an array.
 * Provides two different approaches:
 * 1. Modified Fisher-Yates (Knuth) shuffle - for when array can be modified
 * 2. Reservoir sampling - for streaming data or when array modification is not desired
 *
 * Both algorithms guarantee uniform randomness and handle duplicates correctly.
 */
public class RandomKElements {

    /**
     * Selects k random elements using Modified Fisher-Yates shuffle algorithm.
     * This method creates a copy of the input array to preserve the original.
     *
     * Time Complexity: O(N) - must process full array for copy
     * Space Complexity: O(N) - requires working copy to preserve input
     *
     * @param nums Input array from which to select elements
     * @param k Number of elements to select randomly
     * @return Array containing k randomly selected elements
     * @throws IllegalArgumentException if nums is null, k <= 0, or k > nums.length
     */
    public static int[] getRandomKElements(int[] nums, int k) {
        if (nums == null || k <= 0 || k > nums.length) {
            throw new IllegalArgumentException("Invalid input parameters");
        }

        // Create a copy to avoid modifying the original array
        int[] arr = Arrays.copyOf(nums, nums.length);
        Random rand = new Random();

        // Use Fisher-Yates shuffle algorithm, but only for k elements
        for (int i = 0; i < k; i++) {
            // Generate random index between i and n-1
            int randomIndex = i + rand.nextInt(arr.length - i);

            // Swap elements at positions i and randomIndex
            int temp = arr[i];
            arr[i] = arr[randomIndex];
            arr[randomIndex] = temp;
        }

        // Return first k elements
        return Arrays.copyOf(arr, k);
    }

    /**
     * Selects k random elements using Reservoir Sampling algorithm.
     * This method is particularly useful for:
     * 1. Streaming data where total size is unknown
     * 2. Very large arrays where copying is expensive
     * 3. Cases where modifying original array is not allowed
     *
     * Time Complexity: O(N) - must process each element once
     * Space Complexity: O(k) - only stores k elements
     *
     * @param nums Input array from which to select elements
     * @param k Number of elements to select randomly
     * @return Array containing k randomly selected elements
     * @throws IllegalArgumentException if nums is null, k <= 0, or k > nums.length
     */
    public static int[] getRandomKElementsReservoir(int[] nums, int k) {
        if (nums == null || k <= 0 || k > nums.length) {
            throw new IllegalArgumentException("Invalid input parameters");
        }

        // Initialize reservoir with first k elements
        int[] reservoir = Arrays.copyOf(nums, k);
        Random rand = new Random();

        // Process remaining elements
        for (int i = k; i < nums.length; i++) {
            // Randomly decide whether to include current element
            int j = rand.nextInt(i + 1);

            // If j < k, replace element at index j in reservoir
            if (j < k) {
                reservoir[j] = nums[i];
            }
        }

        return reservoir;
    }

    /**
     * Main method demonstrating both algorithms with example input.
     */
    public static void main(String[] args) {
        int[] nums = {3, 2, 4, 3, 5, 3, 5, 2, 9};
        int k = 4;

        System.out.println("Original array: " + Arrays.toString(nums));
        System.out.println("Random " + k + " elements (Fisher-Yates): " + 
            Arrays.toString(getRandomKElements(nums, k)));
        System.out.println("Random " + k + " elements (Reservoir): " + 
            Arrays.toString(getRandomKElementsReservoir(nums, k)));
    }
}