Combined Find and FindFirst into one function

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

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)
 	}
 }