update the algorithm of removing daughter tracks

Author: wuctlby

PWGHF/D2H/Tasks/taskPtFlucCharmHadrons.cxx

@@ -252,24 +252,31 @@ struct HfTaskPtFlucCharmHadrons {
   template <DecayChannel Channel, typename CandT>
   float removeDaughtersFromMeanPt(const CandT& cand, float rawMeanPt, int n, const std::vector<int>& trkIDs)
   {
-    float meanPt{0.f};
+    float meanPt = rawMeanPt;
+    int removedCount = 0;
+    float removedSumPt = 0.f;
+    auto removeDaug = [&] (int daugID, float daugPt) {
+      if (std::binary_search(trkIDs.begin(), trkIDs.end(), daugID)) {
+        removedSumPt += daugPt;
+        ++removedCount;
+      }
+    };
     if constexpr (Channel == DecayChannel::DplusToPiKPi) {
       std::array<int, 3> daugIDs = {cand.prong0Id(), cand.prong1Id(), cand.prong2Id()};
       std::array<float, 3> daugPts = {cand.ptProng0(), cand.ptProng1(), cand.ptProng2()};
-      for (int iProng = 0; iProng < 3; ++iProng) { // o2-linter: disable=magic-number (for 3-prong)
-        if (std::binary_search(trkIDs.begin(), trkIDs.end(), daugIDs[iProng])) {
-          meanPt = (rawMeanPt * n - daugPts[iProng]) / (n - 1);
-        }
+      for (int iProng = 0; iProng < 3; ++iProng) {
+        removeDaug(daugIDs[iProng], daugPts[iProng]);
       }
     } else if constexpr (Channel == DecayChannel::D0ToPiK || Channel == DecayChannel::D0ToKPi) {
       std::array<int, 2> daugIDs = {cand.prong0Id(), cand.prong1Id()};
       std::array<float, 2> daugPts = {cand.ptProng0(), cand.ptProng1()};
       for (int iProng = 0; iProng < 2; ++iProng) { // o2-linter: disable=magic-number (for 2-prong)
-        if (std::binary_search(trkIDs.begin(), trkIDs.end(), daugIDs[iProng])) {
-          meanPt = (rawMeanPt * n - daugPts[iProng]) / (n - 1);
-        }
+        removeDaug(daugIDs[iProng], daugPts[iProng]);
       }
     }
+    if (removedCount > 0) {
+      meanPt = (rawMeanPt * n - removedSumPt) / (n - removedCount);
+    }
     return meanPt;
   }