File tree Expand file tree Collapse file tree 2 files changed +69
-0
lines changed Expand file tree Collapse file tree 2 files changed +69
-0
lines changed Original file line number Diff line number Diff line change @@ -2444,6 +2444,33 @@ where
2444
2444
}
2445
2445
}
2446
2446
2447
+ /// Remove the `index`th elements along `axis` and shift down elements from higher indexes.
2448
+ ///
2449
+ /// Decreases the length of `axis` by one.
2450
+ ///
2451
+ /// ***Panics** if `axis` or `index` is out of bounds.
2452
+ pub fn shift_remove_index ( & mut self , axis : Axis , index : usize )
2453
+ where
2454
+ S : DataOwned + DataMut ,
2455
+ {
2456
+ // TODO: It's possible to choose to shift the elment to the front or the back here,
2457
+ // whichever is closer.
2458
+ let ( _, mut tail) = self . view_mut ( ) . split_at ( axis, index) ;
2459
+ // shift elements to the back
2460
+ // use swap to keep all elements initialized (as required by owned storage)
2461
+ Zip :: from ( tail. lanes_mut ( axis) ) . for_each ( |mut lane| {
2462
+ let mut lane_iter = lane. iter_mut ( ) ;
2463
+ let mut dst = if let Some ( dst) = lane_iter. next ( ) { dst } else { return } ;
2464
+
2465
+ for elt in lane_iter {
2466
+ std:: mem:: swap ( dst, elt) ;
2467
+ dst = elt;
2468
+ }
2469
+ } ) ;
2470
+ // then slice the axis in place to cut out the removed final element
2471
+ self . slice_axis_inplace ( axis, Slice :: new ( 0 , Some ( -1 ) , 1 ) ) ;
2472
+ }
2473
+
2447
2474
/// Iterates over pairs of consecutive elements along the axis.
2448
2475
///
2449
2476
/// The first argument to the closure is an element, and the second
Original file line number Diff line number Diff line change @@ -2397,3 +2397,45 @@ mod array_cow_tests {
2397
2397
} ) ;
2398
2398
}
2399
2399
}
2400
+
2401
+ #[ test]
2402
+ fn test_shift_remove ( ) {
2403
+ let mut a = arr2 ( & [ [ 1 , 2 , 3 ] ,
2404
+ [ 4 , 5 , 6 ] ,
2405
+ [ 7 , 8 , 9 ] ,
2406
+ [ 10 , 11 , 12 ] ] ) ;
2407
+ a. shift_remove_index ( Axis ( 0 ) , 1 ) ;
2408
+ a. shift_remove_index ( Axis ( 1 ) , 2 ) ;
2409
+ assert_eq ! ( a. shape( ) , & [ 3 , 2 ] ) ;
2410
+ assert_eq ! ( a,
2411
+ array![ [ 1 , 2 ] ,
2412
+ [ 7 , 8 ] ,
2413
+ [ 10 , 11 ] ] ) ;
2414
+
2415
+ let mut a = arr2 ( & [ [ 1 , 2 , 3 ] ,
2416
+ [ 4 , 5 , 6 ] ,
2417
+ [ 7 , 8 , 9 ] ,
2418
+ [ 10 , 11 , 12 ] ] ) ;
2419
+ a. invert_axis ( Axis ( 0 ) ) ;
2420
+ a. shift_remove_index ( Axis ( 0 ) , 1 ) ;
2421
+ a. shift_remove_index ( Axis ( 1 ) , 2 ) ;
2422
+ assert_eq ! ( a. shape( ) , & [ 3 , 2 ] ) ;
2423
+ assert_eq ! ( a,
2424
+ array![ [ 10 , 11 ] ,
2425
+ [ 4 , 5 ] ,
2426
+ [ 1 , 2 ] ] ) ;
2427
+
2428
+ a. shift_remove_index ( Axis ( 1 ) , 1 ) ;
2429
+
2430
+ assert_eq ! ( a. shape( ) , & [ 3 , 1 ] ) ;
2431
+ assert_eq ! ( a,
2432
+ array![ [ 10 ] ,
2433
+ [ 4 ] ,
2434
+ [ 1 ] ] ) ;
2435
+ a. shift_remove_index ( Axis ( 1 ) , 0 ) ;
2436
+ assert_eq ! ( a. shape( ) , & [ 3 , 0 ] ) ;
2437
+ assert_eq ! ( a,
2438
+ array![ [ ] ,
2439
+ [ ] ,
2440
+ [ ] ] ) ;
2441
+ }
You can’t perform that action at this time.
0 commit comments