0
# Indexing and Searching
1

2
Advanced indexing, searching, and selection operations including multi-dimensional indexing, conditional selection, and array searching functions. These operations provide comprehensive tools for accessing, extracting, and manipulating specific array elements based on various criteria.
3

4
## Capabilities
5

6
### Advanced Indexing
7

8
Functions for complex array indexing and element selection.
9

10
```python { .api }
11
def take(a, indices, axis=None, out=None, mode='raise'):
12
    """
13
    Take elements from an array along an axis.
14
    
15
    Parameters:
16
    - a: array_like, input array
17
    - indices: array_like, indices of elements to take
18
    - axis: int, axis over which to select values
19
    - out: cupy.ndarray, optional output array
20
    - mode: str, how to handle out-of-bounds indices
21
    
22
    Returns:
23
    cupy.ndarray: array with selected elements
24
    """
25

26
def take_along_axis(arr, indices, axis):
27
    """
28
    Take values from the input array by matching 1d index and data slices.
29
    
30
    Parameters:
31
    - arr: cupy.ndarray, input array
32
    - indices: cupy.ndarray, indices to take along each 1d slice
33
    - axis: int, axis along which to take 1d slices
34
    
35
    Returns:
36
    cupy.ndarray: array with selected elements
37
    """
38

39
def choose(a, choices, out=None, mode='raise'):
40
    """
41
    Construct an array from an index array and a set of arrays to choose from.
42
    
43
    Parameters:
44
    - a: int array, index array specifying which choice to use
45
    - choices: sequence of arrays, choice arrays
46
    - out: cupy.ndarray, optional output array
47
    - mode: str, how to handle out-of-bounds indices
48
    
49
    Returns:
50
    cupy.ndarray: constructed array
51
    """
52

53
def compress(condition, a, axis=None, out=None):
54
    """
55
    Return selected slices of an array along given axis.
56
    
57
    Parameters:
58
    - condition: 1-D array of bool, which entries to return
59
    - a: array_like, input array
60
    - axis: int, axis along which to take slices
61
    - out: cupy.ndarray, optional output array
62
    
63
    Returns:
64
    cupy.ndarray: compressed array
65
    """
66

67
def extract(condition, arr):
68
    """
69
    Return the elements of an array that satisfy some condition.
70
    
71
    Parameters:
72
    - condition: array_like, condition array
73
    - arr: array_like, input array
74
    
75
    Returns:
76
    cupy.ndarray: 1-D array with elements where condition is True
77
    """
78

79
def select(condlist, choicelist, default=0):
80
    """
81
    Return an array drawn from elements in choicelist, depending on conditions.
82
    
83
    Parameters:
84
    - condlist: list of arrays, conditions determining choice
85
    - choicelist: list of arrays, arrays to choose from
86
    - default: scalar, default value
87
    
88
    Returns:
89
    cupy.ndarray: array with selected elements
90
    """
91
```
92

93
### Index Generation
94

95
Functions for generating index arrays and coordinate arrays.
96

97
```python { .api }
98
def indices(dimensions, dtype=int, sparse=False):
99
    """
100
    Return an array representing the indices of a grid.
101
    
102
    Parameters:
103
    - dimensions: sequence of ints, shape of the grid
104
    - dtype: dtype, data type of result
105
    - sparse: bool, return sparse representation
106
    
107
    Returns:
108
    cupy.ndarray or tuple: grid indices
109
    """
110

111
def ix_(*args):
112
    """
113
    Construct an open mesh from multiple sequences.
114
    
115
    Parameters:
116
    - args: 1-D sequences, input sequences
117
    
118
    Returns:
119
    tuple of cupy.ndarray: open mesh arrays
120
    """
121

122
def unravel_index(indices, shape, order='C'):
123
    """
124
    Converts a flat index or array of flat indices into a tuple of coordinate arrays.
125
    
126
    Parameters:
127
    - indices: array_like, flat indices to convert
128
    - shape: tuple of ints, shape of the array
129
    - order: str, index order ('C' or 'F')
130
    
131
    Returns:
132
    tuple of cupy.ndarray: coordinate arrays
133
    """
134

135
def ravel_multi_index(multi_index, dims, mode='raise', order='C'):
136
    """
137
    Converts a tuple of index arrays into an array of flat indices.
138
    
139
    Parameters:
140
    - multi_index: tuple of array_like, tuple of index arrays
141
    - dims: tuple of ints, shape of array
142
    - mode: str, how to handle out-of-bounds indices
143
    - order: str, index order ('C' or 'F')
144
    
145
    Returns:
146
    cupy.ndarray: flat indices array
147
    """
148

149
def diag_indices(n, ndim=2):
150
    """
151
    Return the indices to access the main diagonal of an array.
152
    
153
    Parameters:
154
    - n: int, size of array for which diagonal indices are returned
155
    - ndim: int, number of dimensions
156
    
157
    Returns:
158
    tuple of cupy.ndarray: diagonal indices
159
    """
160

161
def diag_indices_from(arr):
162
    """
163
    Return the indices to access the main diagonal of an n-dimensional array.
164
    
165
    Parameters:
166
    - arr: array_like, input array
167
    
168
    Returns:
169
    tuple of cupy.ndarray: diagonal indices
170
    """
171
```
172

173
### Array Searching
174

175
Functions for finding elements in arrays.
176

177
```python { .api }
178
def where(condition, x=None, y=None):
179
    """
180
    Return elements chosen from x or y depending on condition.
181
    
182
    Parameters:
183
    - condition: array_like, bool, where True, yield x, otherwise yield y
184
    - x, y: array_like, values from which to choose
185
    
186
    Returns:
187
    cupy.ndarray or tuple: selected elements or indices
188
    """
189

190
def nonzero(a):
191
    """
192
    Return the indices of the elements that are non-zero.
193
    
194
    Parameters:
195
    - a: array_like, input array
196
    
197
    Returns:
198
    tuple of cupy.ndarray: indices of non-zero elements
199
    """
200

201
def flatnonzero(a):
202
    """
203
    Return indices that are non-zero in the flattened version of a.
204
    
205
    Parameters:
206
    - a: array_like, input array
207
    
208
    Returns:
209
    cupy.ndarray: 1-D array of indices
210
    """
211

212
def argwhere(a):
213
    """
214
    Find the indices of array elements that are non-zero, grouped by element.
215
    
216
    Parameters:
217
    - a: array_like, input array
218
    
219
    Returns:
220
    cupy.ndarray: indices of non-zero elements
221
    """
222

223
def searchsorted(a, v, side='left', sorter=None):
224
    """
225
    Find indices where elements should be inserted to maintain order.
226
    
227
    Parameters:
228
    - a: 1-D array_like, sorted input array
229
    - v: array_like, values to insert
230
    - side: str, which side to return insertion point
231
    - sorter: 1-D array_like, optional array of indices sorting a
232
    
233
    Returns:
234
    cupy.ndarray: insertion indices
235
    """
236
```
237

238
### Extrema Finding
239

240
Functions for finding minimum and maximum elements and their indices.
241

242
```python { .api }
243
def argmax(a, axis=None, out=None):
244
    """
245
    Returns the indices of the maximum values along an axis.
246
    
247
    Parameters:
248
    - a: array_like, input array
249
    - axis: int, axis along which to operate
250
    - out: cupy.ndarray, optional output array
251
    
252
    Returns:
253
    cupy.ndarray: indices of maximum values
254
    """
255

256
def argmin(a, axis=None, out=None):
257
    """
258
    Returns the indices of the minimum values along an axis.
259
    
260
    Parameters:
261
    - a: array_like, input array
262
    - axis: int, axis along which to operate
263
    - out: cupy.ndarray, optional output array
264
    
265
    Returns:
266
    cupy.ndarray: indices of minimum values
267
    """
268

269
def nanargmax(a, axis=None):
270
    """
271
    Return the indices of the maximum values in the specified axis ignoring NaNs.
272
    
273
    Parameters:
274
    - a: array_like, input array
275
    - axis: int, axis along which to operate
276
    
277
    Returns:
278
    cupy.ndarray: indices of maximum values
279
    """
280

281
def nanargmin(a, axis=None):
282
    """
283
    Return the indices of the minimum values in the specified axis ignoring NaNs.
284
    
285
    Parameters:
286
    - a: array_like, input array
287
    - axis: int, axis along which to operate
288
    
289
    Returns:
290
    cupy.ndarray: indices of minimum values
291
    """
292
```
293

294
### Array Insertion and Modification
295

296
Functions for modifying arrays at specific indices.
297

298
```python { .api }
299
def place(arr, mask, vals):
300
    """
301
    Change elements of an array based on conditional and input values.
302
    
303
    Parameters:
304
    - arr: cupy.ndarray, array to put data into
305
    - mask: array_like, boolean mask array
306
    - vals: array_like, values to put into arr
307
    """
308

309
def put(a, ind, v, mode='raise'):
310
    """
311
    Replaces specified elements of an array with given values.
312
    
313
    Parameters:
314
    - a: cupy.ndarray, target array
315
    - ind: array_like, target indices
316
    - v: array_like, values to place
317
    - mode: str, how to handle out-of-bounds indices
318
    """
319

320
def putmask(a, mask, values):
321
    """
322
    Changes elements of an array based on conditional and input values.
323
    
324
    Parameters:
325
    - a: cupy.ndarray, target array
326
    - mask: array_like, boolean mask array
327
    - values: array_like, values to assign
328
    """
329

330
def fill_diagonal(a, val, wrap=False):
331
    """
332
    Fill the main diagonal of the given array of any dimensionality.
333
    
334
    Parameters:
335
    - a: cupy.ndarray, array to modify (modified in-place)
336
    - val: scalar, value to be written on the diagonal
337
    - wrap: bool, whether to wrap diagonal for tall matrices
338
    """
339
```
340

341
### Diagonal Operations
342

343
Functions for working with array diagonals.
344

345
```python { .api }
346
def diagonal(a, offset=0, axis1=0, axis2=1):
347
    """
348
    Return specified diagonals.
349
    
350
    Parameters:
351
    - a: array_like, input array
352
    - offset: int, diagonal offset
353
    - axis1: int, first axis
354
    - axis2: int, second axis
355
    
356
    Returns:
357
    cupy.ndarray: diagonal elements
358
    """
359
```
360

361
### Iterator Objects
362

363
Objects for iterating over arrays.
364

365
```python { .api }
366
class flatiter:
367
    """
368
    Flat iterator object to iterate over arrays.
369
    
370
    A flatiter iterator is returned by flat for any array type.
371
    """
372
    def __iter__(self): ...
373
    def __next__(self): ...
374
    def __getitem__(self, key): ...
375
    def __setitem__(self, key, value): ...
376
```
377

378
## Usage Examples
379

380
### Basic Indexing and Selection
381

382
```python
383
import cupy as cp
384

385
# Create test arrays
386
arr = cp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
387
indices = cp.array([0, 2, 1])
388

389
# Take elements along axis
390
result = cp.take(arr, indices, axis=0)
391
print(result)  # [[1 2 3] [7 8 9] [4 5 6]]
392

393
# Extract elements based on condition
394
condition = arr > 5
395
selected = cp.extract(condition, arr)
396
print(selected)  # [6 7 8 9]
397
```
398

399
### Finding Non-zero Elements
400

401
```python
402
import cupy as cp
403

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

407
# Find non-zero elements
408
nonzero_indices = cp.nonzero(arr)
409
print(nonzero_indices)  # (array([0, 1, 1, 2]), array([1, 0, 2, 1]))
410

411
# Get non-zero values
412
nonzero_values = arr[nonzero_indices]
413
print(nonzero_values)  # [1 2 3 4]
414

415
# Alternative using where
416
where_result = cp.where(arr != 0)
417
print(where_result)  # Same as nonzero_indices
418
```
419

420
### Conditional Selection
421

422
```python
423
import cupy as cp
424

425
# Create arrays for conditional selection
426
condition = cp.array([True, False, True, False, True])
427
x = cp.array([1, 2, 3, 4, 5])
428
y = cp.array([10, 20, 30, 40, 50])
429

430
# Select from x where condition is True, y otherwise
431
result = cp.where(condition, x, y)
432
print(result)  # [1 20 3 40 5]
433
```
434

435
### Finding Extrema
436

437
```python
438
import cupy as cp
439

440
# 2D array
441
arr = cp.array([[3, 1, 4], [1, 5, 9], [2, 6, 5]])
442

443
# Find indices of maximum values
444
max_indices = cp.argmax(arr, axis=1)
445
print(max_indices)  # [2 2 1]
446

447
# Find global maximum index (flattened)
448
global_max_idx = cp.argmax(arr)
449
print(global_max_idx)  # 5 (corresponds to value 9)
450

451
# Get the actual maximum values
452
max_values = cp.max(arr, axis=1)
453
print(max_values)  # [4 9 6]
454
```
455

456
### Advanced Indexing with Choose
457

458
```python
459
import cupy as cp
460

461
# Index array
462
indices = cp.array([[0, 1], [2, 0]])
463

464
# Choice arrays
465
choices = [cp.array([[10, 11], [12, 13]]),  # choice 0
466
           cp.array([[20, 21], [22, 23]]),  # choice 1
467
           cp.array([[30, 31], [32, 33]])]  # choice 2
468

469
# Choose elements based on indices
470
result = cp.choose(indices, choices)
471
print(result)  # [[10 21] [32 13]]
472
```
473

474
### Searching in Sorted Arrays
475

476
```python
477
import cupy as cp
478

479
# Sorted array
480
sorted_arr = cp.array([1, 3, 5, 7, 9, 11])
481

482
# Values to insert
483
values = cp.array([2, 6, 8, 12])
484

485
# Find insertion points
486
insertion_points = cp.searchsorted(sorted_arr, values)
487
print(insertion_points)  # [1 3 4 6]
488

489
# Insert values to maintain sorted order
490
result = cp.insert(sorted_arr, insertion_points, values)
491
print(result)  # [1 2 3 5 6 7 8 9 11 12]
492
```