diff --git a/src/sc2lib/sc2_search.cc b/src/sc2lib/sc2_search.cc index 6c9b0fe5..8592017c 100644 --- a/src/sc2lib/sc2_search.cc +++ b/src/sc2lib/sc2_search.cc @@ -103,6 +103,23 @@ std::vector CalculateExpansionLocations(const ObservationInterface* obs } std::vector results = query->Placement(queries); + // Edit the results : allow to build in command structure existing location + Units commandStructures = observation->GetUnits( + [](const Unit& unit) { + return unit.unit_type == UNIT_TYPEID::TERRAN_COMMANDCENTER || unit.unit_type == UNIT_TYPEID::TERRAN_ORBITALCOMMAND || unit.unit_type == UNIT_TYPEID::TERRAN_PLANETARYFORTRESS || + unit.unit_type == UNIT_TYPEID::PROTOSS_NEXUS || + unit.unit_type == UNIT_TYPEID::ZERG_HATCHERY || unit.unit_type == UNIT_TYPEID::ZERG_LAIR || unit.unit_type == UNIT_TYPEID::ZERG_HIVE; + }); + for (auto cc : commandStructures) { + for (int i = 0; i < queries.size(); ++i) { + if (DistanceSquared2D(cc->pos, queries[i].target_pos) < 1.0f) { + results[i] = true; + // we can safely break here, each query is at least one away + break; + } + } + } + size_t start_index = 0; for (int i = 0; i < clusters.size(); ++i) { std::pair >& cluster = clusters[i];