diff --git a/services/companion/java/com/android/server/companion/AssociationStoreImpl.java b/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
index dbcdd0f877a1d..38f8633eaa0a0 100644
--- a/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
+++ b/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
@@ -213,11 +213,9 @@ class AssociationStoreImpl implements AssociationStore {
             final Set<Integer> ids = mAddressMap.get(address);
             if (ids == null) return Collections.emptyList();
 
-            final List<AssociationInfo> associations = new ArrayList<>();
-            for (AssociationInfo association : mIdMap.values()) {
-                if (address.equals(association.getDeviceMacAddress())) {
-                    associations.add(association);
-                }
+            final List<AssociationInfo> associations = new ArrayList<>(ids.size());
+            for (Integer id : ids) {
+                associations.add(mIdMap.get(id));
             }
 
             return Collections.unmodifiableList(associations);