Sorting and Counting

def sort(a, axis=-1, kind=None, order=None):
    """
    Return a sorted copy of an array.
    
    Parameters:
    - a: array_like, input array
    - axis: int, axis along which to sort
    - kind: str, sorting algorithm (ignored, uses optimal GPU algorithm)
    - order: str or list of str, field order for structured arrays
    
    Returns:
    cupy.ndarray: sorted array
    """

def argsort(a, axis=-1, kind=None, order=None):
    """
    Returns the indices that would sort an array.
    
    Parameters:
    - a: array_like, input array
    - axis: int, axis along which to sort
    - kind: str, sorting algorithm (ignored, uses optimal GPU algorithm)
    - order: str or list of str, field order for structured arrays
    
    Returns:
    cupy.ndarray: indices that sort the array
    """

def lexsort(keys, axis=-1):
    """
    Perform an indirect stable sort using a sequence of keys.
    
    Parameters:
    - keys: tuple of array_like, sequence of arrays to use as sort keys
    - axis: int, axis along which to sort
    
    Returns:
    cupy.ndarray: indices for lexicographic sort
    """

def msort(a):
    """
    Return a copy of an array sorted along the first axis.
    
    Parameters:
    - a: array_like, input array
    
    Returns:
    cupy.ndarray: sorted array along first axis
    """

def sort_complex(a):
    """
    Sort a complex array using the real part first, then the imaginary part.
    
    Parameters:
    - a: array_like, input array
    
    Returns:
    cupy.ndarray: sorted complex array
    """

def partition(a, kth, axis=-1, kind=None, order=None):
    """
    Return a partitioned copy of an array.
    
    Parameters:
    - a: array_like, input array
    - kth: int or sequence of ints, element index to partition around
    - axis: int, axis along which to partition
    - kind: str, selection algorithm (ignored)
    - order: str or list of str, field order for structured arrays
    
    Returns:
    cupy.ndarray: partitioned array
    """

def argpartition(a, kth, axis=-1, kind=None, order=None):
    """
    Perform an indirect partition along the given axis.
    
    Parameters:
    - a: array_like, input array
    - kth: int or sequence of ints, element index to partition around
    - axis: int, axis along which to partition
    - kind: str, selection algorithm (ignored)
    - order: str or list of str, field order for structured arrays
    
    Returns:
    cupy.ndarray: indices that partition the array
    """

def count_nonzero(a, axis=None, keepdims=False):
    """
    Counts the number of non-zero values in the array.
    
    Parameters:
    - a: array_like, input array
    - axis: None, int, or tuple, axis along which to count
    - keepdims: bool, whether to keep dimensions in result
    
    Returns:
    int or cupy.ndarray: count of non-zero values
    """

import cupy as cp

# Create unsorted array
arr = cp.array([3, 1, 4, 1, 5, 9, 2, 6])

# Sort array
sorted_arr = cp.sort(arr)
print(sorted_arr)  # [1 1 2 3 4 5 6 9]

# Get sorting indices
sort_indices = cp.argsort(arr)
print(sort_indices)  # [1 3 6 0 2 4 7 5]

# Verify: original array indexed by sort_indices equals sorted array
print(arr[sort_indices])  # [1 1 2 3 4 5 6 9]

import cupy as cp

# 2D array
arr_2d = cp.array([[3, 1, 4], [1, 5, 9], [2, 6, 5]])

# Sort along different axes
sorted_axis0 = cp.sort(arr_2d, axis=0)  # Sort columns
print(sorted_axis0)
# [[1 1 4]
#  [2 5 5] 
#  [3 6 9]]

sorted_axis1 = cp.sort(arr_2d, axis=1)  # Sort rows
print(sorted_axis1)
# [[1 3 4]
#  [1 5 9]
#  [2 5 6]]

import cupy as cp

# Multiple sort keys
surnames = cp.array([3, 1, 2, 1, 3])    # Primary key
firstnames = cp.array([1, 2, 1, 1, 2])  # Secondary key

# Sort by surname first, then by firstname
indices = cp.lexsort((firstnames, surnames))
print(indices)  # [3 1 2 0 4]

# Result is sorted first by surnames, then by firstnames
print(surnames[indices])   # [1 1 2 3 3]
print(firstnames[indices]) # [1 2 1 1 2]

import cupy as cp

# Array to partition
arr = cp.array([7, 1, 9, 3, 5, 2, 8, 4, 6])

# Partition around the 4th smallest element (index 4 when sorted)
partitioned = cp.partition(arr, 4)
print(partitioned)  # Elements <= 5th smallest on left, rest on right

# Get partition indices
partition_indices = cp.argpartition(arr, 4)
print(partition_indices)

# The 4th index contains the 5th smallest element
kth_element = partitioned[4]
print(f"5th smallest element: {kth_element}")

import cupy as cp

# Array with some zeros
arr = cp.array([[0, 1, 2], [3, 0, 4], [5, 6, 0]])

# Count non-zero elements
total_nonzero = cp.count_nonzero(arr)
print(total_nonzero)  # 6

# Count along axes
nonzero_per_row = cp.count_nonzero(arr, axis=1)
print(nonzero_per_row)  # [2 2 2]

nonzero_per_col = cp.count_nonzero(arr, axis=0)
print(nonzero_per_col)  # [2 2 2]

import cupy as cp

# Complex array
complex_arr = cp.array([1+2j, 3+1j, 2+3j, 1+1j])

# Sort complex numbers (by real part first, then imaginary)
sorted_complex = cp.sort_complex(complex_arr)
print(sorted_complex)  # [1+1j 1+2j 2+3j 3+1j]

import cupy as cp

# Student scores in multiple subjects
math_scores = cp.array([85, 92, 78, 96, 88])
english_scores = cp.array([88, 85, 92, 89, 91])
student_ids = cp.array([101, 102, 103, 104, 105])

# Sort students by total score (descending)
total_scores = math_scores + english_scores
sort_indices = cp.argsort(-total_scores)  # Negative for descending

print("Ranking by total score:")
for i, idx in enumerate(sort_indices):
    print(f"{i+1}. Student {student_ids[idx]}: "
          f"Math={math_scores[idx]}, English={english_scores[idx]}, "
          f"Total={total_scores[idx]}")