Statistics and Sorting

def mean(a, axis=None, dtype=None, out=None, keepdims=False):
    """
    Compute arithmetic mean along specified axis.

    Parameters:
    - a: input array
    - axis: axis or axes along which to compute mean
    - dtype: data type for computation
    - out: output array
    - keepdims: keep dimensions of original array

    Returns:
    cupy.ndarray: arithmetic mean
    """

def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
    """
    Compute standard deviation along specified axis.

    Parameters:
    - a: input array
    - axis: axis or axes along which to compute std
    - dtype: data type for computation
    - out: output array
    - ddof: degrees of freedom correction
    - keepdims: keep dimensions

    Returns:
    cupy.ndarray: standard deviation
    """

def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
    """Compute variance along specified axis."""

def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
    """Compute median along specified axis."""

def average(a, axis=None, weights=None, returned=False):
    """
    Compute weighted average along specified axis.

    Parameters:
    - a: input array
    - axis: axis along which to average
    - weights: weights for averaging
    - returned: return sum of weights

    Returns:
    cupy.ndarray or tuple: weighted average
    """

def nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
    """Compute mean ignoring NaNs."""

def nanstd(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
    """Compute standard deviation ignoring NaNs."""

def nanvar(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
    """Compute variance ignoring NaNs."""

def nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False):
    """Compute median ignoring NaNs."""

def amin(a, axis=None, out=None, keepdims=False, initial=None, where=True):
    """Return minimum values along axis."""

def amax(a, axis=None, out=None, keepdims=False, initial=None, where=True):
    """Return maximum values along axis."""

def min(a, axis=None, out=None, keepdims=False, initial=None, where=True):
    """Alias for amin."""

def max(a, axis=None, out=None, keepdims=False, initial=None, where=True):
    """Alias for amax."""

def nanmin(a, axis=None, out=None, keepdims=False):
    """Return minimum values ignoring NaNs."""

def nanmax(a, axis=None, out=None, keepdims=False):
    """Return maximum values ignoring NaNs."""

def ptp(a, axis=None, out=None, keepdims=False):
    """
    Range of values (maximum - minimum) along axis.

    Parameters:
    - a: input array
    - axis: axis along which to compute range
    - out: output array
    - keepdims: keep dimensions

    Returns:
    cupy.ndarray: peak-to-peak values
    """

def percentile(a, q, axis=None, out=None, overwrite_input=False, 
              method='linear', keepdims=False):
    """
    Compute qth percentile along specified axis.

    Parameters:
    - a: input array
    - q: percentile(s) to compute
    - axis: axis along which to compute percentiles
    - out: output array
    - overwrite_input: allow input modification
    - method: interpolation method
    - keepdims: keep dimensions

    Returns:
    cupy.ndarray: qth percentiles
    """

def quantile(a, q, axis=None, out=None, overwrite_input=False,
            method='linear', keepdims=False):
    """Compute quantiles along specified axis."""

def corrcoef(x, y=None, rowvar=True, bias=None, ddof=None):
    """
    Return Pearson correlation coefficients.

    Parameters:
    - x: input array
    - y: additional input array
    - rowvar: whether rows represent variables
    - bias: bias correction (deprecated)
    - ddof: degrees of freedom (deprecated)

    Returns:
    cupy.ndarray: correlation coefficient matrix
    """

def cov(m, y=None, rowvar=True, bias=False, ddof=None, fweights=None, aweights=None):
    """
    Estimate covariance matrix.

    Parameters:
    - m: input array
    - y: additional input array
    - rowvar: whether rows represent variables
    - bias: use biased estimator
    - ddof: degrees of freedom correction
    - fweights: frequency weights
    - aweights: analytic weights

    Returns:
    cupy.ndarray: covariance matrix
    """

def correlate(a, v, mode='valid'):
    """
    Cross-correlation of two 1-dimensional sequences.

    Parameters:
    - a: first input sequence
    - v: second input sequence
    - mode: output size ('full', 'valid', 'same')

    Returns:
    cupy.ndarray: cross-correlation
    """

def histogram(a, bins=10, range=None, normed=None, weights=None, density=None):
    """
    Compute histogram of a set of data.

    Parameters:
    - a: input data
    - bins: number of bins or bin edges
    - range: lower and upper range of bins
    - normed: normalize histogram (deprecated)
    - weights: weights for each value
    - density: normalize to create probability density

    Returns:
    tuple: (hist, bin_edges)
    """

def histogram2d(x, y, bins=10, range=None, normed=None, weights=None, density=None):
    """
    Compute 2D histogram of two data samples.

    Parameters:
    - x, y: input data arrays
    - bins: number of bins or bin edges
    - range: array of ranges for each dimension
    - normed: normalize histogram (deprecated)
    - weights: weights for each sample
    - density: normalize to create probability density

    Returns:
    tuple: (H, xedges, yedges)
    """

def histogramdd(sample, bins=10, range=None, normed=None, weights=None, density=None):
    """
    Compute multidimensional histogram.

    Parameters:
    - sample: input data array
    - bins: number of bins for each dimension
    - range: sequence of ranges for each dimension
    - normed: normalize histogram (deprecated)
    - weights: weights for each sample
    - density: normalize to create probability density

    Returns:
    tuple: (H, edges)
    """

def bincount(x, weights=None, minlength=0):
    """
    Count occurrences of each value in array of non-negative ints.

    Parameters:
    - x: input array of non-negative integers
    - weights: weights for each value
    - minlength: minimum number of bins

    Returns:
    cupy.ndarray: counts for each value
    """

def digitize(x, bins, right=False):
    """
    Return indices of bins to which each value belongs.

    Parameters:
    - x: input array
    - bins: array of bins
    - right: whether intervals include right edge

    Returns:
    cupy.ndarray: bin indices
    """

def sort(a, axis=-1, kind=None, order=None):
    """
    Return sorted copy of array.

    Parameters:
    - a: input array
    - axis: axis along which to sort
    - kind: sorting algorithm (ignored, uses merge sort)
    - order: field order for structured arrays

    Returns:
    cupy.ndarray: sorted array
    """

def argsort(a, axis=-1, kind=None, order=None):
    """
    Return indices that would sort array.

    Parameters:
    - a: input array
    - axis: axis along which to sort
    - kind: sorting algorithm
    - order: field order for structured arrays

    Returns:
    cupy.ndarray: indices for sorted array
    """

def lexsort(keys, axis=-1):
    """
    Perform indirect stable sort using multiple keys.

    Parameters:
    - keys: sequence of arrays to use as sort keys
    - axis: axis along which to sort

    Returns:
    cupy.ndarray: indices for lexicographically sorted array
    """

def msort(a):
    """
    Return sorted copy along first axis.

    Parameters:
    - a: input array

    Returns:
    cupy.ndarray: sorted array
    """

def sort_complex(a):
    """
    Sort complex array using real part first, then imaginary part.

    Parameters:
    - a: input complex array

    Returns:
    cupy.ndarray: sorted complex array
    """

def partition(a, kth, axis=-1, kind='introselect', order=None):
    """
    Return partitioned copy where kth element is in correct position.

    Parameters:
    - a: input array
    - kth: element index for partitioning
    - axis: axis along which to partition
    - kind: partitioning algorithm
    - order: field order for structured arrays

    Returns:
    cupy.ndarray: partitioned array
    """

def argpartition(a, kth, axis=-1, kind='introselect', order=None):
    """Return indices that would partition array."""

def argmax(a, axis=None, out=None):
    """
    Return indices of maximum values along axis.

    Parameters:
    - a: input array
    - axis: axis along which to search
    - out: output array

    Returns:
    cupy.ndarray: indices of maximum values
    """

def argmin(a, axis=None, out=None):
    """Return indices of minimum values along axis."""

def nanargmax(a, axis=None):
    """Return indices of maximum values ignoring NaNs."""

def nanargmin(a, axis=None):
    """Return indices of minimum values ignoring NaNs."""

def argwhere(a):
    """
    Find indices of array elements that are non-zero.

    Parameters:
    - a: input array

    Returns:
    cupy.ndarray: indices of non-zero elements
    """

def nonzero(a):
    """
    Return indices of elements that are non-zero.

    Parameters:
    - a: input array

    Returns:
    tuple: arrays of indices
    """

def flatnonzero(a):
    """Return indices of flattened array that are non-zero."""

def where(condition, x=None, y=None):
    """
    Return elements chosen from x or y depending on condition.

    Parameters:
    - condition: boolean array
    - x: values where condition is True
    - y: values where condition is False

    Returns:
    cupy.ndarray: array with elements from x or y
    """

def searchsorted(a, v, side='left', sorter=None):
    """
    Find indices where elements should be inserted to maintain order.

    Parameters:
    - a: sorted input array
    - v: values to insert
    - side: insertion side ('left' or 'right')
    - sorter: array of indices that sort a

    Returns:
    cupy.ndarray: insertion indices
    """

def count_nonzero(a, axis=None, keepdims=False):
    """
    Count number of non-zero values in array.

    Parameters:
    - a: input array
    - axis: axis or axes to count along
    - keepdims: keep dimensions of original array

    Returns:
    int or cupy.ndarray: count of non-zero values
    """

import cupy as cp

# Create sample data
data = cp.random.normal(10, 2, (1000, 50))

# Compute basic statistics
mean_val = cp.mean(data)
std_val = cp.std(data)
var_val = cp.var(data)
median_val = cp.median(data)

print(f"Mean: {mean_val:.4f}")
print(f"Std: {std_val:.4f}")
print(f"Variance: {var_val:.4f}")
print(f"Median: {median_val:.4f}")

# Statistics along specific axis
row_means = cp.mean(data, axis=1)  # Mean of each row
col_stds = cp.std(data, axis=0)    # Std of each column

print(f"Row means shape: {row_means.shape}")
print(f"Column stds shape: {col_stds.shape}")

import cupy as cp

# Create test data
data = cp.random.random((100, 100))

# Find min/max values
min_val = cp.min(data)
max_val = cp.max(data)
range_val = cp.ptp(data)  # peak-to-peak

print(f"Min: {min_val:.4f}")
print(f"Max: {max_val:.4f}")
print(f"Range: {range_val:.4f}")

# Percentiles
percentiles = cp.percentile(data, [25, 50, 75, 90, 95])
print(f"Percentiles (25,50,75,90,95): {percentiles}")

# Quantiles (same as percentiles but with 0-1 scale)
quantiles = cp.quantile(data, [0.25, 0.5, 0.75])
print(f"Quantiles (0.25,0.5,0.75): {quantiles}")

import cupy as cp

# Create data with NaN values
data = cp.random.random((100, 100))
data[cp.random.random((100, 100)) < 0.1] = cp.nan  # 10% NaN values

# Regular statistics (will return NaN if any NaN present)
regular_mean = cp.mean(data)
regular_std = cp.std(data)

# NaN-aware statistics
nan_mean = cp.nanmean(data)
nan_std = cp.nanstd(data)
nan_min = cp.nanmin(data)
nan_max = cp.nanmax(data)

print(f"Regular mean: {regular_mean}")
print(f"NaN-aware mean: {nan_mean:.4f}")
print(f"NaN-aware std: {nan_std:.4f}")
print(f"NaN-aware range: {nan_min:.4f} to {nan_max:.4f}")

import cupy as cp

# Create correlated data
n_samples = 1000
x = cp.random.normal(0, 1, n_samples)
y = 0.8 * x + 0.6 * cp.random.normal(0, 1, n_samples)  # Correlated with x
z = cp.random.normal(0, 1, n_samples)  # Independent

# Stack into matrix (each row is a variable)
data = cp.stack([x, y, z])

# Compute correlation matrix
corr_matrix = cp.corrcoef(data)
print("Correlation matrix:")
print(corr_matrix)

# Compute covariance matrix
cov_matrix = cp.cov(data)
print("\nCovariance matrix:")
print(cov_matrix)

# Cross-correlation of two sequences
x_seq = cp.random.random(100)
y_seq = cp.random.random(100)
cross_corr = cp.correlate(x_seq, y_seq, mode='full')
print(f"\nCross-correlation shape: {cross_corr.shape}")

import cupy as cp

# Create sample data
data = cp.random.normal(0, 1, 10000)

# 1D histogram
hist, bin_edges = cp.histogram(data, bins=50, range=(-4, 4))
print(f"Histogram shape: {hist.shape}")
print(f"Bin edges shape: {bin_edges.shape}")

# Weighted histogram
weights = cp.random.random(len(data))
weighted_hist, _ = cp.histogram(data, bins=50, weights=weights)

# 2D histogram
x = cp.random.normal(0, 1, 5000)
y = cp.random.normal(0, 1, 5000)
hist_2d, x_edges, y_edges = cp.histogram2d(x, y, bins=30)
print(f"2D histogram shape: {hist_2d.shape}")

# Count occurrences
integers = cp.random.randint(0, 10, 1000)
counts = cp.bincount(integers)
print(f"Counts: {counts}")

# Digitize continuous data
bin_indices = cp.digitize(data, bins=cp.linspace(-3, 3, 10))
print(f"Bin indices range: {cp.min(bin_indices)} to {cp.max(bin_indices)}")

import cupy as cp

# Create unsorted data
data = cp.random.random((5, 10))

# Sort array
sorted_data = cp.sort(data, axis=1)  # Sort each row
print("Original data (first row):")
print(data[0])
print("Sorted data (first row):")
print(sorted_data[0])

# Get sorting indices
sort_indices = cp.argsort(data, axis=1)
print("Sort indices (first row):")
print(sort_indices[0])

# Verify sorting
reconstructed = data[0, sort_indices[0]]
print("Reconstructed (should match sorted):")
print(reconstructed)

# Multi-dimensional sort
data_3d = cp.random.random((10, 20, 30))
sorted_3d = cp.sort(data_3d, axis=2)  # Sort along last axis

import cupy as cp

# Lexicographic sorting
# Sort by multiple keys (e.g., sort by y first, then by x)
x = cp.array([1, 3, 2, 1, 3, 2])
y = cp.array([3, 1, 2, 1, 3, 1])

# Sort by y first, then x (note order: primary key last)
lex_indices = cp.lexsort([x, y])
print("Lexsort indices:", lex_indices)
print("x sorted:", x[lex_indices])
print("y sorted:", y[lex_indices])

# Partial sorting (partition)
large_array = cp.random.random(1000)
k = 100  # Find 100 smallest elements

# Partition so that k smallest elements are in first k positions
partitioned = cp.partition(large_array, k)
print(f"100th smallest element: {partitioned[k-1]}")
print(f"Verification - max of first 100: {cp.max(partitioned[:k])}")
print(f"Verification - min of last 900: {cp.min(partitioned[k:])}")

import cupy as cp

# Create test data
data = cp.random.random((50, 50))

# Find locations of extreme values
max_pos = cp.argmax(data)
min_pos = cp.argmin(data)

# Convert flat indices to 2D coordinates
max_coords = cp.unravel_index(max_pos, data.shape)
min_coords = cp.unravel_index(min_pos, data.shape)

print(f"Max value {cp.max(data):.4f} at position {max_coords}")
print(f"Min value {cp.min(data):.4f} at position {min_coords}")

# Find all positions above threshold
threshold = 0.9
high_positions = cp.argwhere(data > threshold)
print(f"Found {len(high_positions)} positions above {threshold}")

# Search in sorted array
sorted_array = cp.sort(cp.random.random(1000))
values_to_find = cp.array([0.1, 0.5, 0.9])
insertion_points = cp.searchsorted(sorted_array, values_to_find)
print(f"Insertion points: {insertion_points}")

# Count non-zero elements
sparse_data = cp.random.random((100, 100))
sparse_data[sparse_data < 0.9] = 0  # Make 90% zeros
nonzero_count = cp.count_nonzero(sparse_data)
print(f"Non-zero elements: {nonzero_count} out of {sparse_data.size}")

import cupy as cp
import numpy as np
import time

# Large dataset for performance testing
n = 10**7
data_gpu = cp.random.random(n)
data_cpu = cp.asnumpy(data_gpu)

# GPU sorting
start = time.time()
sorted_gpu = cp.sort(data_gpu)
cp.cuda.Device().synchronize()
gpu_time = time.time() - start

# CPU sorting
start = time.time()
sorted_cpu = np.sort(data_cpu)
cpu_time = time.time() - start

print(f"GPU sort time: {gpu_time:.4f}s")
print(f"CPU sort time: {cpu_time:.4f}s")
print(f"Speedup: {cpu_time/gpu_time:.2f}x")

# Verify correctness
gpu_result_cpu = cp.asnumpy(sorted_gpu)
max_diff = np.max(np.abs(gpu_result_cpu - sorted_cpu))
print(f"Max difference: {max_diff}")