Indexing and Selection

def take(a, indices, axis=None, out=None, mode='raise'):
    """Take elements from array along specified axis.
    
    Args:
        a: Input array
        indices: Indices of elements to take
        axis: Axis along which to take elements
        out: Output array
        mode: How to handle out-of-bounds indices ('raise', 'wrap', 'clip')
        
    Returns:
        cupy.ndarray: Selected elements
    """

def take_along_axis(arr, indices, axis):
    """Take values from input array by matching 1D index array.
    
    Args:
        arr: Input array
        indices: Indices array with same ndim as arr
        axis: Axis along which to take values
        
    Returns:
        cupy.ndarray: Selected values
    """

def compress(condition, a, axis=None, out=None):
    """Select slices along axis using boolean condition.
    
    Args:
        condition: Boolean 1-D array selecting slices
        a: Input array
        axis: Axis along which to select
        out: Output array
        
    Returns:
        cupy.ndarray: Compressed array
    """

def diagonal(a, offset=0, axis1=0, axis2=1):
    """Extract diagonal elements from 2-D array.
    
    Args:
        a: Input array (≥2D)
        offset: Diagonal offset (0 for main diagonal)
        axis1: First axis for diagonal
        axis2: Second axis for diagonal
        
    Returns:
        cupy.ndarray: Diagonal elements
    """

def diag(v, k=0):
    """Extract diagonal or construct diagonal array.
    
    Args:
        v: Input array or 1-D array for diagonal construction
        k: Diagonal offset
        
    Returns:
        cupy.ndarray: Diagonal elements or diagonal matrix
    """

def choose(a, choices, out=None, mode='raise'):
    """Use index array to construct new array from choice arrays.
    
    Args:
        a: Index array (integers)
        choices: Sequence of arrays to choose from
        out: Output array
        mode: How to handle out-of-bounds indices
        
    Returns:
        cupy.ndarray: Array constructed from choices
    """

def select(condlist, choicelist, default=0):
    """Select elements from choicelist based on conditions.
    
    Args:
        condlist: List of boolean conditions
        choicelist: List of arrays/scalars to choose from
        default: Default value when no condition is True
        
    Returns:
        cupy.ndarray: Array with selected elements
    """

def extract(condition, arr):
    """Extract elements satisfying condition.
    
    Args:
        condition: Boolean array or condition
        arr: Input array
        
    Returns:
        cupy.ndarray: 1-D array of elements where condition is True
    """

def where(condition, x=None, y=None):
    """Select elements from x or y based on condition.
    
    Args:
        condition: Boolean array condition
        x: Array/scalar for True elements
        y: Array/scalar for False elements
        
    Returns:
        cupy.ndarray or tuple: Selected elements or indices
    """

def indices(dimensions, dtype=int, sparse=False):
    """Return grid of indices for given dimensions.
    
    Args:
        dimensions: Shape of output grid
        dtype: Data type for indices
        sparse: Return sparse representation
        
    Returns:
        cupy.ndarray: Grid indices array
    """

def ix_(*args):
    """Construct open mesh from multiple sequences.
    
    Args:
        *args: 1-D sequences for each dimension
        
    Returns:
        tuple: Tuple of arrays for open mesh indexing
    """

def ravel_multi_index(multi_index, dims, mode='raise', order='C'):
    """Convert multi-dimensional index to flat index.
    
    Args:
        multi_index: Tuple of index arrays
        dims: Shape of array being indexed
        mode: How to handle out-of-bounds indices
        order: Memory layout ('C' or 'F')
        
    Returns:
        cupy.ndarray: Flat indices
    """

def unravel_index(indices, shape, order='C'):
    """Convert flat index to multi-dimensional index.
    
    Args:
        indices: Flat index array
        shape: Shape to unravel into
        order: Memory layout ('C' or 'F')
        
    Returns:
        tuple: Tuple of index arrays
    """

def nonzero(a):
    """Return indices of non-zero elements.
    
    Args:
        a: Input array
        
    Returns:
        tuple: Tuple of index arrays for each dimension
    """

def flatnonzero(a):
    """Return indices of non-zero elements in flattened array.
    
    Args:
        a: Input array
        
    Returns:
        cupy.ndarray: 1-D array of flat indices
    """

def argwhere(a):
    """Return indices of non-zero elements grouped by element.
    
    Args:
        a: Input array
        
    Returns:
        cupy.ndarray: 2-D array where each row is an index
    """

def argmax(a, axis=None, out=None, keepdims=False):
    """Indices of maximum values along axis.
    
    Args:
        a: Input array
        axis: Axis for computation
        out: Output array
        keepdims: Keep reduced dimensions
        
    Returns:
        cupy.ndarray: Indices of maximum elements
    """

def argmin(a, axis=None, out=None, keepdims=False):
    """Indices of minimum values along axis.
    
    Args:
        a: Input array
        axis: Axis for computation
        out: Output array
        keepdims: Keep reduced dimensions
        
    Returns:
        cupy.ndarray: Indices of minimum elements
    """

def nanargmax(a, axis=None, out=None, keepdims=False):
    """Indices of maximum values ignoring NaNs.
    
    Args:
        a: Input array
        axis: Axis for computation
        out: Output array
        keepdims: Keep reduced dimensions
        
    Returns:
        cupy.ndarray: Indices of maximum elements (ignoring NaN)
    """

def nanargmin(a, axis=None, out=None, keepdims=False):
    """Indices of minimum values ignoring NaNs.
    
    Args:
        a: Input array
        axis: Axis for computation
        out: Output array
        keepdims: Keep reduced dimensions
        
    Returns:
        cupy.ndarray: Indices of minimum elements (ignoring NaN)
    """

def argsort(a, axis=-1, kind=None, order=None):
    """Indices that would sort array along axis.
    
    Args:
        a: Input array
        axis: Axis to sort along
        kind: Sorting algorithm (None, 'quicksort', 'mergesort', 'heapsort')
        order: Field order for structured arrays
        
    Returns:
        cupy.ndarray: Array of indices that sort input
    """

def argpartition(a, kth, axis=-1, kind='introselect', order=None):
    """Indices that partition array around kth element.
    
    Args:
        a: Input array
        kth: Element index around which to partition
        axis: Axis to partition along
        kind: Partitioning algorithm
        order: Field order for structured arrays
        
    Returns:
        cupy.ndarray: Array of partition indices
    """

def searchsorted(a, v, side='left', sorter=None):
    """Find indices to insert elements to maintain sorted order.
    
    Args:
        a: Input sorted array
        v: Values to insert
        side: Insert position ('left' or 'right')
        sorter: Optional array of indices that sort a
        
    Returns:
        cupy.ndarray: Insertion indices
    """

def put(a, ind, v, mode='raise'):
    """Put values into array at specified indices.
    
    Args:
        a: Target array (modified in-place)
        ind: Target indices
        v: Values to insert
        mode: How to handle out-of-bounds indices
    """

def place(arr, mask, vals):
    """Place values into array based on boolean mask.
    
    Args:
        arr: Target array (modified in-place)
        mask: Boolean array indicating positions
        vals: Values to place
    """

def putmask(a, mask, values):
    """Put values where mask is True.
    
    Args:
        a: Target array (modified in-place)
        mask: Boolean mask
        values: Replacement values
    """

def fill_diagonal(a, val, wrap=False):
    """Fill main diagonal of array.
    
    Args:
        a: Array to modify (≥2D)
        val: Value to fill diagonal with
        wrap: Wrap diagonal for non-square arrays
    """

def diag_indices(n, ndim=2):
    """Return indices for diagonal elements.
    
    Args:
        n: Size of arrays for which indices are returned
        ndim: Number of dimensions
        
    Returns:
        tuple: Indices for diagonal elements
    """

def diag_indices_from(arr):
    """Return indices for diagonal of n-dimensional array.
    
    Args:
        arr: Input array
        
    Returns:
        tuple: Indices for diagonal elements
    """

def mask_indices(n, mask_func, k=0):
    """Return indices for masking based on function.
    
    Args:
        n: Number of rows/columns in output
        mask_func: Function defining mask pattern
        k: Optional parameter for mask function
        
    Returns:
        tuple: Indices satisfying mask condition
    """

def tril_indices(n, k=0, m=None):
    """Return indices for lower triangle of array.
    
    Args:
        n: Number of rows in output
        k: Diagonal offset
        m: Number of columns (defaults to n)
        
    Returns:
        tuple: Lower triangle indices
    """

def tril_indices_from(arr, k=0):
    """Return indices for lower triangle of given array.
    
    Args:
        arr: Input array
        k: Diagonal offset
        
    Returns:
        tuple: Lower triangle indices
    """

def triu_indices(n, k=0, m=None):
    """Return indices for upper triangle of array.
    
    Args:
        n: Number of rows in output
        k: Diagonal offset
        m: Number of columns (defaults to n)
        
    Returns:
        tuple: Upper triangle indices
    """

def triu_indices_from(arr, k=0):
    """Return indices for upper triangle of given array.
    
    Args:
        arr: Input array
        k: Diagonal offset
        
    Returns:
        tuple: Upper triangle indices
    """

def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
    """Create coordinate arrays from coordinate vectors.
    
    Args:
        *xi: 1-D arrays for coordinates
        copy: Return copies of input arrays
        sparse: Return sparse grid
        indexing: Cartesian ('xy') or matrix ('ij') indexing
        
    Returns:
        list: Coordinate arrays
    """

class flatiter:
    """Iterator for flattened array.
    
    Provides iteration over array as if it were 1-dimensional.
    
    Attributes:
        base: Base array being iterated
        coords: Current coordinates
        index: Current flat index
    """
    
    def __iter__(self):
        """Return iterator object."""
        pass
        
    def __next__(self):
        """Get next element."""
        pass
        
    def __getitem__(self, key):
        """Get element(s) by index/slice."""
        pass
        
    def __setitem__(self, key, value):
        """Set element(s) by index/slice.""" 
        pass

class nditer:
    """Multi-dimensional iterator object.
    
    Args:
        op: Array(s) to iterate over
        flags: Iteration flags
        op_flags: Per-operand flags
        op_dtypes: Per-operand data types
        order: Iteration order
        casting: Casting rule
        op_axes: Per-operand axes
        itershape: Iteration shape
        buffersize: Buffer size
        
    Provides advanced iteration patterns with broadcasting,
    buffering, and multi-array iteration.
    """
    
    def __init__(self, op, flags=None, op_flags=None, op_dtypes=None,
                 order='K', casting='safe', op_axes=None, itershape=None,
                 buffersize=0):
        pass

import cupy as cp

# Create sample array
data = cp.arange(24).reshape(4, 6)
print(f"Original array:\n{data}")

# Basic indexing
element = data[1, 3]                    # Single element
row = data[2]                           # Entire row
column = data[:, 1]                     # Entire column  
subarray = data[1:3, 2:5]              # Rectangular slice

print(f"Element [1,3]: {element}")
print(f"Row 2: {row}")
print(f"Column 1: {column}")
print(f"Subarray [1:3, 2:5]:\n{subarray}")

# Take elements by indices
indices = cp.array([0, 2, 1, 3])
taken_rows = cp.take(data, indices, axis=0)
taken_elements = cp.take(data.flatten(), cp.array([5, 10, 15, 20]))

print(f"Taken rows (indices {indices}):\n{taken_rows}")
print(f"Taken elements: {taken_elements}")

# Diagonal extraction
diag_main = cp.diagonal(data)           # Main diagonal
diag_upper = cp.diagonal(data, offset=1)  # Upper diagonal
diag_lower = cp.diagonal(data, offset=-1) # Lower diagonal

print(f"Main diagonal: {diag_main}")
print(f"Upper diagonal: {diag_upper}")
print(f"Lower diagonal: {diag_lower}")

import cupy as cp

# Create test data
temperatures = cp.array([18.5, 22.1, 28.7, 31.4, 19.8, 25.3, 33.1, 16.9])
cities = cp.array(['NYC', 'LA', 'MIA', 'PHX', 'SEA', 'CHI', 'LAS', 'BOS'])

# Boolean indexing
hot_weather = temperatures > 25
cold_weather = temperatures < 20
moderate_weather = (temperatures >= 20) & (temperatures <= 30)

hot_cities = cities[hot_weather]
cold_cities = cities[cold_weather]
hot_temps = temperatures[hot_weather]

print(f"Hot cities (>25°C): {hot_cities}")
print(f"Hot temperatures: {hot_temps}")
print(f"Cold cities (<20°C): {cold_cities}")

# Fancy indexing with arrays
indices = cp.array([1, 3, 5, 7])       # Select specific indices
selected_temps = temperatures[indices]
selected_cities = cities[indices]

print(f"Selected temperatures: {selected_temps}")
print(f"Selected cities: {selected_cities}")

# Multi-dimensional boolean indexing
matrix = cp.random.randn(5, 8)
positive = matrix > 0
negative = matrix < 0
extreme = cp.abs(matrix) > 1.5

print(f"Matrix shape: {matrix.shape}")
print(f"Positive elements: {cp.sum(positive)}")
print(f"Negative elements: {cp.sum(negative)}")
print(f"Extreme values: {cp.sum(extreme)}")

# Extract extreme values and their positions
extreme_values = matrix[extreme]
extreme_positions = cp.argwhere(extreme)

print(f"First 5 extreme values: {extreme_values[:5]}")
print(f"First 5 extreme positions:\n{extreme_positions[:5]}")

import cupy as cp

# Sample data - student scores
math_scores = cp.array([85, 92, 76, 88, 91, 73, 95, 82])
english_scores = cp.array([78, 89, 94, 82, 86, 88, 91, 79])
science_scores = cp.array([90, 87, 81, 95, 89, 76, 93, 85])

# Use where for conditional selection
# Select higher of math or english score for each student
higher_score = cp.where(math_scores > english_scores, math_scores, english_scores)

print(f"Math scores:    {math_scores}")
print(f"English scores: {english_scores}")
print(f"Higher scores:  {higher_score}")

# Multi-condition selection with where
# Assign letter grades based on best score
best_scores = cp.maximum(cp.maximum(math_scores, english_scores), science_scores)
grades = cp.where(best_scores >= 90, 'A',
                 cp.where(best_scores >= 80, 'B',
                         cp.where(best_scores >= 70, 'C', 'F')))

print(f"Best scores: {best_scores}")
print(f"Grades: {grades}")

# Using choose for more complex selection
# Choose subject based on student preference (0=math, 1=english, 2=science)
preferences = cp.array([0, 2, 1, 0, 2, 1, 0, 1])  # Student preferences
score_arrays = cp.array([math_scores, english_scores, science_scores])
preferred_scores = cp.choose(preferences, score_arrays)

print(f"Preferences: {preferences}")
print(f"Preferred scores: {preferred_scores}")

# Using select for multiple conditions
conditions = [
    best_scores >= 95,  # Excellent
    best_scores >= 85,  # Good  
    best_scores >= 75,  # Fair
]
choices = ['Excellent', 'Good', 'Fair']
performance = cp.select(conditions, choices, default='Poor')

print(f"Performance ratings: {performance}")

import cupy as cp

# Create sample dataset
data = cp.random.randn(1000, 10)
outlier_threshold = 2.5

# Find indices of extremes
max_indices = cp.argmax(data, axis=1)       # Max in each row
min_indices = cp.argmin(data, axis=1)       # Min in each row
overall_max_idx = cp.argmax(data)           # Overall max (flat index)
overall_min_idx = cp.argmin(data)           # Overall min (flat index)

print(f"Data shape: {data.shape}")
print(f"Overall max at flat index: {overall_max_idx}")
print(f"Overall min at flat index: {overall_min_idx}")

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

print(f"Max at coordinates: {max_coords}")
print(f"Min at coordinates: {min_coords}")
print(f"Max value: {data[max_coords]:.3f}")
print(f"Min value: {data[min_coords]:.3f}")

# Find outliers (values beyond threshold)
outliers_mask = cp.abs(data) > outlier_threshold
outlier_positions = cp.argwhere(outliers_mask)
outlier_values = data[outliers_mask]

print(f"Number of outliers: {len(outlier_values)}")
print(f"First 5 outlier positions:\n{outlier_positions[:5]}")
print(f"First 5 outlier values: {outlier_values[:5]}")

# Non-zero elements
nonzero_mask = data != 0
nonzero_indices = cp.nonzero(nonzero_mask)
nonzero_count = len(nonzero_indices[0])

print(f"Non-zero elements: {nonzero_count}/{data.size}")

# Extract elements satisfying complex conditions
condition = (data > 0) & (cp.abs(data) < 1.0)  # Positive values between 0 and 1
extracted = cp.extract(condition, data)

print(f"Values between 0 and 1: {len(extracted)}")
print(f"Mean of extracted values: {cp.mean(extracted):.3f}")

import cupy as cp

# Create unsorted data
scores = cp.array([87, 92, 76, 88, 91, 73, 95, 82, 89, 78])
names = cp.array(['Alice', 'Bob', 'Charlie', 'David', 'Eve', 
                  'Frank', 'Grace', 'Henry', 'Iris', 'Jack'])

print(f"Original scores: {scores}")
print(f"Original names: {names}")

# Get indices that would sort the scores
sort_indices = cp.argsort(scores)
reverse_sort_indices = cp.argsort(-scores)  # Descending order

sorted_scores = scores[sort_indices]
sorted_names = names[sort_indices]
top_scores = scores[reverse_sort_indices]
top_names = names[reverse_sort_indices]

print(f"Sorted scores (ascending): {sorted_scores}")
print(f"Sorted names (by score): {sorted_names}")
print(f"Top scores (descending): {top_scores}")
print(f"Top names (by score): {top_names}")

# Partial sorting - find top 3 students
top_k = 3
top_k_indices = cp.argpartition(-scores, top_k)[:top_k]
top_k_scores = scores[top_k_indices]
top_k_names = names[top_k_indices]

# Sort the top k for proper ranking
top_k_order = cp.argsort(-top_k_scores)
final_top_scores = top_k_scores[top_k_order]
final_top_names = top_k_names[top_k_order]

print(f"Top {top_k} scores: {final_top_scores}")
print(f"Top {top_k} names: {final_top_names}")

# Searching in sorted array
target_scores = cp.array([80, 90, 95])
sorted_scores = cp.sort(scores)

# Find insertion positions
insert_left = cp.searchsorted(sorted_scores, target_scores, side='left')
insert_right = cp.searchsorted(sorted_scores, target_scores, side='right')

print(f"Sorted scores: {sorted_scores}")
print(f"Target scores: {target_scores}")
print(f"Insert positions (left): {insert_left}")
print(f"Insert positions (right): {insert_right}")

# Count how many scores are below each target
below_count = insert_left
print(f"Scores below targets: {below_count}")

import cupy as cp

# Create coordinate grids
x = cp.linspace(-2, 2, 5)
y = cp.linspace(-1, 1, 4)

# Create meshgrid for function evaluation
X, Y = cp.meshgrid(x, y)
print(f"X coordinates:\n{X}")
print(f"Y coordinates:\n{Y}")

# Evaluate function over grid
Z = X**2 + Y**2  # Distance from origin squared
print(f"Function values:\n{Z}")

# Advanced indexing with meshgrid
# Select elements where distance is less than 2
distance_mask = Z < 2.0
selected_X = X[distance_mask]
selected_Y = Y[distance_mask] 
selected_Z = Z[distance_mask]

print(f"Selected coordinates (distance < 2):")
print(f"X: {selected_X}")
print(f"Y: {selected_Y}")
print(f"Z: {selected_Z}")

# Open mesh indexing with ix_
rows = cp.array([0, 2])
cols = cp.array([1, 3, 4])
row_idx, col_idx = cp.ix_(rows, cols)

print(f"Open mesh indices:")
print(f"Row indices:\n{row_idx}")
print(f"Col indices:\n{col_idx}")

# Use open mesh to index into array
data = cp.arange(20).reshape(4, 5)
selected = data[row_idx, col_idx]
print(f"Original data:\n{data}")
print(f"Selected subarray:\n{selected}")

# Multi-dimensional index operations
shape = (6, 8)
flat_indices = cp.array([5, 12, 25, 33, 41])
multi_indices = cp.unravel_index(flat_indices, shape)

print(f"Flat indices: {flat_indices}")
print(f"Multi-dimensional indices:")
print(f"  Row indices: {multi_indices[0]}")
print(f"  Col indices: {multi_indices[1]}")

# Convert back to flat indices
recovered_flat = cp.ravel_multi_index(multi_indices, shape)
print(f"Recovered flat indices: {recovered_flat}")

import cupy as cp

# Create sample array
data = cp.random.randn(6, 8)
print(f"Original data shape: {data.shape}")
print(f"Original data (first 3 rows):\n{data[:3]}")

# Boolean masking for modification
# Replace negative values with zero
negative_mask = data < 0
data_clipped = data.copy()
data_clipped[negative_mask] = 0

print(f"Negative elements: {cp.sum(negative_mask)}")
print(f"After clipping negatives (first 3 rows):\n{data_clipped[:3]}")

# Use place for conditional replacement
data_replaced = data.copy()
outlier_mask = cp.abs(data) > 2.0
cp.place(data_replaced, outlier_mask, cp.nan)

print(f"Outlier elements: {cp.sum(outlier_mask)}")
print(f"Outliers replaced with NaN: {cp.sum(cp.isnan(data_replaced))}")

# Use putmask for different replacement strategy
data_putmask = data.copy()
extreme_mask = cp.abs(data) > 1.5
replacement_values = cp.sign(data) * 1.5  # Cap at ±1.5
cp.putmask(data_putmask, extreme_mask, replacement_values[extreme_mask])

print(f"Values capped at ±1.5")
print(f"Max absolute value: {cp.max(cp.abs(data_putmask)):.3f}")

# Direct indexing for complex modifications
# Set border elements to specific value
border_value = -999
data_border = data.copy()
data_border[0, :] = border_value      # Top row
data_border[-1, :] = border_value     # Bottom row
data_border[:, 0] = border_value      # Left column
data_border[:, -1] = border_value     # Right column

print(f"Border elements set to {border_value}")
print(f"Modified array:\n{data_border}")

# Diagonal operations
square_data = data[:6, :6]  # Make it square
cp.fill_diagonal(square_data, 0)
diag_indices = cp.diag_indices(6)
square_data[diag_indices] = cp.arange(6) * 10  # Set diagonal to multiples of 10

print(f"Modified diagonal:\n{square_data}")

# Triangle operations
upper_tri_indices = cp.triu_indices(6, k=1)
lower_tri_indices = cp.tril_indices(6, k=-1)

tri_data = cp.zeros((6, 6))
tri_data[upper_tri_indices] = 1    # Upper triangle
tri_data[lower_tri_indices] = -1   # Lower triangle

print(f"Triangle pattern:\n{tri_data}")

import cupy as cp
import time

# Large dataset for performance testing
n = 1_000_000
large_data = cp.random.randn(n)
indices = cp.random.randint(0, n, size=100_000)

print(f"Performance test with {n:,} elements")

# Time different indexing methods
methods = [
    ('Direct indexing', lambda: large_data[indices]),
    ('Take method', lambda: cp.take(large_data, indices)),
    ('Boolean masking', lambda: large_data[large_data > 0]),
    ('Where condition', lambda: cp.where(large_data > 0, large_data, 0)),
    ('Compress', lambda: cp.compress(large_data > 0, large_data)),
]

for name, method in methods:
    start_time = time.perf_counter()
    result = method()
    cp.cuda.Stream.null.synchronize()  # Wait for GPU
    end_time = time.perf_counter()
    
    print(f"{name:16}: {(end_time - start_time)*1000:6.2f} ms, "
          f"result size: {result.size:8,}")

# Memory-efficient indexing patterns
# Process large array in chunks to avoid memory issues
chunk_size = 100_000
results = []

print(f"\nChunked processing:")
for i in range(0, n, chunk_size):
    chunk_end = min(i + chunk_size, n)
    chunk = large_data[i:chunk_end]
    
    # Process chunk (example: find outliers)
    outliers = chunk[cp.abs(chunk) > 2.0]
    if len(outliers) > 0:
        results.append(outliers)
    
    if i // chunk_size < 5:  # Show first few chunks
        print(f"Chunk {i//chunk_size + 1}: "
              f"processed {chunk_end - i:6,} elements, "
              f"found {len(outliers):3} outliers")

if results:
    all_outliers = cp.concatenate(results)
    print(f"Total outliers found: {len(all_outliers)}")
else:
    print("No outliers found")