Use map indexing to speed up PickByKeys and OmitByKeys

samber · Apr 25, 2024 · b590c5b · b590c5b
1 parent 71d8341
commit b590c5b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/map.go b/map.go
@@ -49,8 +49,8 @@ func PickBy[K comparable, V any](in map[K]V, predicate func(key K, value V) bool
 // Play: https://go.dev/play/p/R1imbuci9qU
 func PickByKeys[K comparable, V any](in map[K]V, keys []K) map[K]V {
 	r := map[K]V{}
-	for k, v := range in {
-		if Contains(keys, k) {
+	for _, k := range keys {
+		if v, ok := in[k]; ok {
 			r[k] = v
 		}
 	}
@@ -86,9 +86,10 @@ func OmitBy[K comparable, V any](in map[K]V, predicate func(key K, value V) bool
 func OmitByKeys[K comparable, V any](in map[K]V, keys []K) map[K]V {
 	r := map[K]V{}
 	for k, v := range in {
-		if !Contains(keys, k) {
-			r[k] = v
-		}
+		r[k] = v
+	}
+	for _, k := range keys {
+		delete(r, k)
 	}
 	return r
 }

diff --git a/map_test.go b/map_test.go
@@ -55,7 +55,7 @@ func TestPickByKeys(t *testing.T) {
 	t.Parallel()
 	is := assert.New(t)
 
-	r1 := PickByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz"})
+	r1 := PickByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz", "qux"})
 
 	is.Equal(r1, map[string]int{"foo": 1, "baz": 3})
 }
@@ -84,7 +84,7 @@ func TestOmitByKeys(t *testing.T) {
 	t.Parallel()
 	is := assert.New(t)
 
-	r1 := OmitByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz"})
+	r1 := OmitByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz", "qux"})
 
 	is.Equal(r1, map[string]int{"bar": 2})
 }