diff --git a/README.md b/README.md index f9ccde3..e0f2657 100644 --- a/README.md +++ b/README.md @@ -747,6 +747,18 @@ l := lo.DropRightWhile([]string{"a", "aa", "aaa", "aa", "aa"}, func(val string) [[play](https://go.dev/play/p/3-n71oEC0Hz)] +### DropByIndex + +Drops elements from a slice or array by the index. A negative index will drop elements from the end of the slice. + +```go +l := lo.Drop([]int{0, 1, 2, 3, 4, 5}, 2, 4, -1) +// []int{2, 3} +``` + +[[play](https://go.dev/play/p/JswS7vXRJP2)] + + ### Reject The opposite of Filter, this method returns the elements of collection that predicate does not return truthy for. diff --git a/slice.go b/slice.go index 5d7b479..5e8c155 100644 --- a/slice.go +++ b/slice.go @@ -1,6 +1,7 @@ package lo import ( + "sort" "math/rand" "github.com/samber/lo/internal/constraints" @@ -417,6 +418,38 @@ func DropRightWhile[T any, Slice ~[]T](collection Slice, predicate func(item T) return append(result, collection[:i+1]...) } +// DropByIndex drops elements from a slice or array by the index. +// A negative index will drop elements from the end of the slice. +// Play: https://go.dev/play/p/bPIH4npZRxS +func DropByIndex[T any](collection []T, indexes ...int) []T { + initialSize := len(collection) + if initialSize == 0 { + return make([]T, 0) + } + + for i := range indexes { + if indexes[i] < 0 { + indexes[i] = initialSize + indexes[i] + } + } + + indexes = Uniq(indexes) + sort.Ints(indexes) + + result := make([]T, 0, initialSize) + result = append(result, collection...) + + for i := range indexes { + if indexes[i]-i < 0 || indexes[i]-i >= initialSize-i { + continue + } + + result = append(result[:indexes[i]-i], result[indexes[i]-i+1:]...) + } + + return result +} + // Reject is the opposite of Filter, this method returns the elements of collection that predicate does not return truthy for. // Play: https://go.dev/play/p/YkLMODy1WEL func Reject[T any, Slice ~[]T](collection Slice, predicate func(item T, index int) bool) Slice { diff --git a/slice_benchmark_test.go b/slice_benchmark_test.go index 6b08511..7a1ff5d 100644 --- a/slice_benchmark_test.go +++ b/slice_benchmark_test.go @@ -151,6 +151,26 @@ func BenchmarkDropRightWhile(b *testing.B) { } } +func BenchmarkDropByIndex(b *testing.B) { + for _, n := range lengths { + strs := genSliceString(n) + b.Run(fmt.Sprintf("strings_%d", n), func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = DropByIndex(strs, n/4) + } + }) + } + + for _, n := range lengths { + ints := genSliceInt(n) + b.Run(fmt.Sprintf("ints%d", n), func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = DropByIndex(ints, n/4) + } + }) + } +} + func BenchmarkReplace(b *testing.B) { lengths := []int{1_000, 10_000, 100_000} for _, n := range lengths { diff --git a/slice_example_test.go b/slice_example_test.go index b5b49d2..1ee90ce 100644 --- a/slice_example_test.go +++ b/slice_example_test.go @@ -296,6 +296,15 @@ func ExampleDropRightWhile() { // Output: [0 1 2] } +func ExampleDropByIndex() { + list := []int{0, 1, 2, 3, 4, 5} + + result := DropByIndex(list, 2) + + fmt.Printf("%v", result) + // Output: [0 1 3 4 5] +} + func ExampleReject() { list := []int{0, 1, 2, 3, 4, 5} diff --git a/slice_test.go b/slice_test.go index cd39b9c..0917e17 100644 --- a/slice_test.go +++ b/slice_test.go @@ -571,6 +571,27 @@ func TestDropRightWhile(t *testing.T) { is.IsType(nonempty, allStrings, "type preserved") } +func TestDropByIndex(t *testing.T) { + t.Parallel() + is := assert.New(t) + + is.Equal([]int{1, 2, 3, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, 0)) + is.Equal([]int{3, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, 0, 1, 2)) + is.Equal([]int{0, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, -4, -2, -3)) + is.Equal([]int{0, 2, 3, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, -4, -4)) + is.Equal([]int{2, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, 3, 1, 0)) + is.Equal([]int{0, 1, 3, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, 2)) + is.Equal([]int{0, 1, 2, 3}, DropByIndex([]int{0, 1, 2, 3, 4}, 4)) + is.Equal([]int{0, 1, 2, 3, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, 5)) + is.Equal([]int{0, 1, 2, 3, 4}, DropByIndex([]int{0, 1, 2, 3, 4}, 100)) + is.Equal([]int{0, 1, 2, 3}, DropByIndex([]int{0, 1, 2, 3, 4}, -1)) + is.Equal([]int{}, DropByIndex([]int{}, 0, 1)) + is.Equal([]int{}, DropByIndex([]int{42}, 0, 1)) + is.Equal([]int{}, DropByIndex([]int{42}, 1, 0)) + is.Equal([]int{}, DropByIndex([]int{}, 1)) + is.Equal([]int{}, DropByIndex([]int{1}, 0)) +} + func TestReject(t *testing.T) { t.Parallel() is := assert.New(t)