[mathcore] Fix TKDTreeBinning::FindBin returning out-of-range bins

TKDTreeBinning computed fNBins independently from the kd-tree it builds.
SetNBins picks a bucket size of fDataSize / nBins (integer division), but
the underlying TKDTree creates ceil(fNPoints / bucketSize) terminal nodes.
When the requested number of bins does not divide the data size evenly, the
bucket size is rounded down and the tree ends up with more terminal nodes
than fNBins. Since FindBin returns a terminal-node index, it could return a
value >= GetNBins(), pointing at a bin with no edges and no content:

```c++
binning.GetNBins()                         # 1001
binning.FindBin([1.2, 0.5, 0.5, 0.5, 0.5]) # 1004
binning.GetBinContent(1004)                # RuntimeWarning: No such bin
```

Set fNBins to the actual number of terminal nodes (GetNNodes() + 1) after
building the tree, so FindBin always returns an index within
[0, GetNBins()) and every bin has valid edges and content. This also
self-corrects existing objects on read-back, as the Streamer rebuilds the
tree via SetNBins.

Also fix SetBinsContent to read the per-node point counts directly from the
tree (GetNPointsNode) instead of an incorrect last-bin heuristic, so bin
contents are exact and sum to the data size.

Add a regression test covering the non-evenly-dividing case.

Closes #10786.

🤖 Done with the help of AI.

Author: guitargeek

math/mathcore/src/TKDTreeBinning.cxx

@@ -136,6 +136,12 @@ void TKDTreeBinning::SetNBins(UInt_t bins) {
          fDataBins = new TKDTreeID(fDataSize, fDim, fDataSize / (fNBins - remainingData)); // TKDTree input is data size, data dimension and the content size of bins ("bucket" size)
          SetTreeData();
          fDataBins->Build();
+         // The kd-tree determines the actual number of bins: every terminal node
+         // is a bin and there are GetNNodes() + 1 of them. When the requested bin
+         // count does not divide the data size evenly, the bucket size is rounded
+         // down and the tree ends up with more terminal nodes than the requested
+         // number of bins. We must therefore use the real number of terminal nodes.
+         fNBins = fDataBins->GetNNodes() + 1;
          SetBinsEdges();
          SetBinsContent();
       } else {
@@ -238,12 +244,13 @@ void TKDTreeBinning::SetTreeData() {
 }
 
 void TKDTreeBinning::SetBinsContent() {
-   // Sets the bins' content
+   // Sets the bins' content from the actual number of points in each terminal
+   // node of the kd-tree. All terminal nodes hold fBucketSize points except
+   // possibly the last one, so query the tree directly instead of guessing.
    fBinsContent.resize(fNBins);
+   UInt_t nTreeNodes = fDataBins->GetNNodes();
    for (UInt_t i = 0; i < fNBins; ++i)
-      fBinsContent[i] = fDataBins->GetBucketSize();
-   if ( fDataSize % fNBins != 0 )
-      fBinsContent[fNBins - 1] = fDataSize % (fNBins-1);
+      fBinsContent[i] = fDataBins->GetNPointsNode(i + nTreeNodes);
 }
 
 void TKDTreeBinning::SetBinsEdges() {

math/mathcore/test/testkdTreeBinning.cxx

@@ -169,6 +169,66 @@ int testkdTreeBinning() {
 
 }
 
+// Regression test for the case where the requested number of bins does not
+// divide the data size evenly. In that situation the kd-tree builds more
+// terminal nodes (bins) than naively expected, and FindBin() must never return
+// an index outside [0, GetNBins()). See
+// https://github.com/root-project/root/issues/10784 about
+// TKDTreeBinning::FindBin returning non-existent bins. Returns the number of
+// detected failures (0 on success) so that the caller can turn it into a
+// non-zero process exit code: ROOT's Error() only prints, it does not by
+// itself make the test fail under ctest.
+int testkdTreeBinningFindBinRange() {
+
+   int nfail = 0;
+
+   const UInt_t DATASZ  = 100500;  // deliberately NOT a multiple of NBINS
+   const UInt_t DATADIM = 5;
+   const UInt_t NBINS   = 1000;
+
+   std::vector<Double_t> smp(DATASZ * DATADIM);
+   TRandom3 r;
+   r.SetSeed(1);
+   for (UInt_t i = 0; i < DATADIM; ++i)
+      for (UInt_t j = 0; j < DATASZ; ++j)
+         smp[DATASZ * i + j] = r.Uniform(-1., 1.);
+
+   TKDTreeBinning kdBins(DATASZ, DATADIM, smp, NBINS);
+
+   const UInt_t nbins = kdBins.GetNBins();
+
+   // The number of bins must match the number of terminal nodes of the kd-tree.
+   if ((int)nbins != kdBins.GetTree()->GetNNodes() + 1) {
+      Error("testkdTreeBinningFindBinRange", "GetNBins() (%u) != number of kd-tree terminal nodes (%d)",
+            nbins, kdBins.GetTree()->GetNNodes() + 1);
+      ++nfail;
+   }
+
+   // Every data point must be assigned to a valid bin.
+   std::vector<Double_t> point(DATADIM);
+   for (UInt_t j = 0; j < DATASZ; ++j) {
+      for (UInt_t i = 0; i < DATADIM; ++i)
+         point[i] = smp[DATASZ * i + j];
+      UInt_t bin = kdBins.FindBin(point.data());
+      if (bin >= nbins) {
+         Error("testkdTreeBinningFindBinRange", "FindBin returned out-of-range bin %u (NBins = %u)", bin, nbins);
+         ++nfail;
+         break;
+      }
+   }
+
+   // The total bin content must add up to the data size.
+   Long64_t total = 0;
+   for (UInt_t i = 0; i < nbins; ++i)
+      total += kdBins.GetBinContent(i);
+   if (total != (Long64_t)DATASZ) {
+      Error("testkdTreeBinningFindBinRange", "Sum of bin contents (%lld) != data size (%u)", total, DATASZ);
+      ++nfail;
+   }
+
+   return nfail;
+}
+
 int main(int argc, char **argv)
 {
   // Parse command line arguments
@@ -197,6 +257,8 @@ int main(int argc, char **argv)
 
    int nfail = testkdTreeBinning();
 
+   nfail += testkdTreeBinningFindBinRange();
+
    if ( showGraphics )
    {
       theApp->Run();