diff --git a/floats.go b/floats.go index 22ac16d6..432bbe52 100644 --- a/floats.go +++ b/floats.go @@ -5,7 +5,7 @@ import "math" // InsufficientElements is an error type used by FindFirst type InsufficientElements struct{} -func (i InsufficientElements) Error() string { +func (i *InsufficientElements) Error() string { return "Insufficient elements found" } @@ -109,38 +109,51 @@ func EqLen(slices ...[]float64) bool { return true } -// Find applies a function returning a boolean to the elements of the slice -// and returns a list of indices for which the value is true -func Find(s []float64, f func(float64) bool) (inds []int) { - // Not sure what an appropriate capacity is here. Don't want to make - // it the length of the slice because if the slice is large that is - // a lot of potentially wasted memory - inds = make([]int, 0) - for i, val := range s { - if f(val) { - inds = append(inds, i) - } - } - return inds -} +// Find finds the first k indices of the slice s for which +// the function f returns true and stores them in the slice +// inds. If k < 0, all such elements are found. +// Find will reslice inds to have 0 length, and will append +// found indices to inds. +// If there are fewer than k elements in s satisfying f, +// all of the found elements will be returned along with an +// InsufficientElements error +// TODO: Add example for nil slice, appending to the end of a larger slice +// reusing memory, insufficient elements +func Find(inds []int, k int, s []float64, f func(float64) bool) ([]int, error) { -// FindFirst applies a function returning a boolean to the elements of the slice -// and returns a list of the first k indices for which the value is true. -// If there are fewer than k indices for which the value is true, it returns -// the found indices and an error. -func FindFirst(s []float64, f func(float64) bool, k int) (inds []int, err error) { - count := 0 - inds = make([]int, 0, k) + // inds is also returned to allow for calling with nil + + // Reslice inds to have zero length + inds = inds[:0] + + // If zero elements requested, can just return + if k == 0 { + return inds, nil + } + + // If k < 0, return all of the found indices + if k < 0 { + for i, val := range s { + if f(val) { + inds = append(inds, i) + } + } + return inds, nil + } + + // Otherwise, find the first k elements + nFound := 0 for i, val := range s { if f(val) { inds = append(inds, i) - count++ - if count == k { + nFound++ + if nFound == k { return inds, nil } } } - return inds, InsufficientElements{} + // Finished iterating over the loop, which means k elements were not found + return inds, &InsufficientElements{} } // LogSpan returns a set of N equally spaced points in log space between l and u, where N diff --git a/floats_test.go b/floats_test.go index 52435e90..58348c09 100644 --- a/floats_test.go +++ b/floats_test.go @@ -1,9 +1,9 @@ package sliceops import ( - "fmt" "math" "math/rand" + "strconv" "testing" ) @@ -111,50 +111,63 @@ func TestEqLen(t *testing.T) { } } +func eqIntSlice(one, two []int) string { + if len(one) != len(two) { + return "Length mismatch" + } + for i, val := range one { + if val != two[i] { + return "Index " + strconv.Itoa(i) + " mismatch" + } + } + return "" +} + func TestFind(t *testing.T) { s := []float64{3, 4, 1, 7, 5} f := func(v float64) bool { return v > 3.5 } - trueInds := []int{1, 3, 4} - inds := Find(s, f) - if len(inds) != len(trueInds) { - t.Errorf("Wrong number of elements returned") - return - } - for i, val := range trueInds { - if inds[i] != val { - t.Errorf("Index mismatch") - fmt.Println(trueInds) - fmt.Println(inds) - return - } - } -} + allTrueInds := []int{1, 3, 4} -func TestFindFirst(t *testing.T) { - s := []float64{3, 4, 1, 7, 5} - f := func(v float64) bool { return v > 3.5 } - trueInds := []int{1, 3} - k := 2 - inds, err := FindFirst(s, f, k) + // Test finding first two elements + inds, err := Find(nil, 2, s, f) if err != nil { - t.Errorf("Incorrectly did not find enough elements") + t.Errorf("Find first two: Improper error return") } - if len(inds) != len(trueInds) { - t.Errorf("Wrong number of elements returned") - return + trueInds := allTrueInds[:2] + str := eqIntSlice(inds, trueInds) + if str != "" { + t.Errorf("Find first two: " + str) } - for i, val := range trueInds { - if inds[i] != val { - t.Errorf("Index mismatch") - fmt.Println(trueInds) - fmt.Println(inds) - return - } + + // Test finding first two elements with non nil slice + inds = []int{1, 2, 3, 4, 5, 6} + inds, err = Find(inds, 2, s, f) + if err != nil { + t.Errorf("Find first two non-nil: Improper error return") } - f = func(v float64) bool { return v > 6.0 } - inds, err = FindFirst(s, f, k) + str = eqIntSlice(inds, trueInds) + if str != "" { + t.Errorf("Find first two non-nil: " + str) + } + + // Test finding too many elements + inds, err = Find(inds, 4, s, f) if err == nil { - t.Errorf("Incorrectly found enough elements") + t.Errorf("Request too many: No error returned") + } + str = eqIntSlice(inds, allTrueInds) + if str != "" { + t.Errorf("Request too many: Does not match all of the inds: " + str) + } + + // Test finding all elements + inds, err = Find(nil, -1, s, f) + if err != nil { + t.Errorf("Find all: Improper error returned") + } + str = eqIntSlice(inds, allTrueInds) + if str != "" { + t.Errorf("Find all: Does not match all of the inds: " + str) } }