diff --git a/src/wifi/model/ap-wifi-mac.cc b/src/wifi/model/ap-wifi-mac.cc
index df5e2ca7f..cd354ea16 100644
--- a/src/wifi/model/ap-wifi-mac.cc
+++ b/src/wifi/model/ap-wifi-mac.cc
@@ -107,6 +107,7 @@ ApWifiMac::~ApWifiMac ()
 {
   NS_LOG_FUNCTION (this);
   m_staList.clear ();
+  m_addressIdMap.clear ();
   m_nonErpStations.clear ();
   m_nonHtStations.clear ();
 }
@@ -746,6 +747,7 @@ ApWifiMac::SendAssocResp (Mac48Address to, bool success, bool isReassoc)
         {
           aid = GetNextAssociationId ();
           m_staList.insert (std::make_pair (aid, to));
+          m_addressIdMap.insert (std::make_pair (to, aid));
         }
       assoc.SetAssociationId (aid);
     }
@@ -1359,6 +1361,7 @@ ApWifiMac::Receive (Ptr<WifiMacQueueItem> mpdu)
                   if (j->second == from)
                     {
                       m_staList.erase (j);
+                      m_addressIdMap.erase (from);
                       break;
                     }
                 }
@@ -1465,14 +1468,13 @@ ApWifiMac::GetStaList (void) const
 uint16_t
 ApWifiMac::GetAssociationId (Mac48Address addr) const
 {
-  for (auto & it : m_staList)
+  auto it = m_addressIdMap.find (addr);
+
+  if (it == m_addressIdMap.end ())
     {
-      if (it.second == addr)
-        {
-          return it.first;
-        }
+      return SU_STA_ID;
     }
-  return SU_STA_ID;
+  return it->second;
 }
 
 uint8_t
diff --git a/src/wifi/model/ap-wifi-mac.h b/src/wifi/model/ap-wifi-mac.h
index 3dfbf6f5b..d7e719b85 100644
--- a/src/wifi/model/ap-wifi-mac.h
+++ b/src/wifi/model/ap-wifi-mac.h
@@ -307,6 +307,8 @@ private:
   Ptr<UniformRandomVariable> m_beaconJitter; //!< UniformRandomVariable used to randomize the time of the first beacon
   bool m_enableBeaconJitter;                 //!< Flag whether the first beacon should be generated at random time
   std::map<uint16_t, Mac48Address> m_staList; //!< Map of all stations currently associated to the AP with their association ID
+  //!< Maps MAC addresses of associated stations to their association ID
+  std::unordered_map<Mac48Address, uint16_t, WifiAddressHash> m_addressIdMap;
   std::list<Mac48Address> m_nonErpStations;  //!< List of all non-ERP stations currently associated to the AP
   std::list<Mac48Address> m_nonHtStations;   //!< List of all non-HT stations currently associated to the AP
   bool m_enableNonErpProtection;             //!< Flag whether protection mechanism is used or not when non-ERP STAs are present within the BSS