go-hep.org/x/hep@v0.38.1/sliceop/sliceop.go (about)

     1  // Copyright ©2021 The go-hep Authors.  All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build go1.18
     6  
     7  // Package sliceop provides operations on slices not available in the stdlib
     8  // slices package.
     9  package sliceop // import "go-hep.org/x/hep/sliceop"
    10  
    11  import "errors"
    12  
    13  var (
    14  	errLength           = errors.New("sliceop: length mismatch")
    15  	errSortedIndices    = errors.New("sliceop: indices not sorted")
    16  	errDuplicateIndices = errors.New("sliceop: duplicate indices")
    17  )
    18  
    19  // Filter creates a slice with all the elements x_i of src for which f(x_i) is true.
    20  // Filter uses dst as work buffer, storing elements at the start of the slice.
    21  // Filter clears dst if a slice is passed, and allocates a new slice if dst is nil.
    22  func Filter[S ~[]E, E any](dst, src S, f func(v E) bool) S {
    23  
    24  	if dst == nil {
    25  		dst = make(S, 0, len(src))
    26  	}
    27  
    28  	dst = dst[:0]
    29  	for _, x := range src {
    30  		if f(x) {
    31  			dst = append(dst, x)
    32  		}
    33  	}
    34  
    35  	return dst
    36  }
    37  
    38  // Map creates a slice with all the elements f(x_i) where x_i are elements from src.
    39  // Map uses dst as work buffer, storing elements at the start of the slice.
    40  // Map allocates a new slice if dst is nil.
    41  // Map will panic if the lengths of src and dst differ.
    42  func Map[T, U any](dst []U, src []T, f func(v T) U) []U {
    43  
    44  	if dst == nil {
    45  		dst = make([]U, len(src))
    46  	}
    47  
    48  	if len(src) != len(dst) {
    49  		panic(errLength)
    50  	}
    51  
    52  	for i, x := range src {
    53  		dst[i] = f(x)
    54  	}
    55  	return dst
    56  }
    57  
    58  // Find creates a slice with all indices corresponding to elements for which f(x) is true.
    59  // Find uses dst as work buffer, storing indices at the start of the slice.
    60  // Find clears dst if a slice is passed, and allocates a new slice if dst is nil.
    61  func Find[S ~[]E, E any](dst []int, src S, f func(v E) bool) []int {
    62  
    63  	if dst == nil {
    64  		dst = make([]int, 0, len(src))
    65  	}
    66  
    67  	dst = dst[:0]
    68  	for i, x := range src {
    69  		if f(x) {
    70  			dst = append(dst, i)
    71  		}
    72  	}
    73  
    74  	return dst
    75  }
    76  
    77  // Take creates a sub-slice of src with all elements indiced by the provided indices.
    78  // Take uses dst as work buffer, storing elements at the start of the slice.
    79  // Take clears dst if a slice is passed, and allocates a new slice if dst is nil.
    80  // Take will panic if indices is not sorted or has duplicates.
    81  // Take will panic if length of indices is larger than length of src.
    82  // Take will panic if length of indices is different from length of dst.
    83  func Take[S ~[]E, E any](dst, src S, indices []int) S {
    84  
    85  	if len(indices) > len(src) {
    86  		panic(errLength)
    87  	}
    88  
    89  	if dst == nil {
    90  		dst = make(S, len(indices))
    91  	}
    92  
    93  	if len(dst) != len(indices) {
    94  		panic(errLength)
    95  	}
    96  
    97  	if len(indices) == 0 {
    98  		return dst
    99  	}
   100  
   101  	dst[0] = src[indices[0]]
   102  	var (
   103  		v0 = indices[0]
   104  		nn = len(indices)
   105  	)
   106  	for i := 1; i < nn; i++ {
   107  		v1 := indices[i]
   108  		switch {
   109  		case v0 < v1:
   110  			// ok.
   111  		case v0 == v1:
   112  			panic(errDuplicateIndices)
   113  		case v0 > v1:
   114  			panic(errSortedIndices)
   115  		}
   116  		dst[i-1] = src[v0]
   117  		v0 = v1
   118  	}
   119  	dst[nn-1] = src[v0]
   120  
   121  	return dst
   122  }
   123  
   124  // Resize returns a slice of size n, reusing the storage of the provided
   125  // slice, appending new elements if the capacity is not sufficient.
   126  func Resize[S ~[]E, E any](sli S, n int) S {
   127  	if m := cap(sli); m < n {
   128  		sli = sli[:m]
   129  		sli = append(sli, make(S, n-m)...)
   130  	}
   131  	sli = sli[:n]
   132  	return sli
   133  }