Repository: sebp/scikit-survival (language: python)

Issue:
RandomSurvivalForest unusally slow
I attempted to fit a RandomSurvivalForest with 500,000 training instances and 140 features on a machine with 90GB memory. Unfortunately, after **hours** of waiting, the program ran out of memory and crashed. I wasn't able to see any progress even though the parameter 'verbose' was set to 1.

However, I was able to fit a sklearn.RandomForestRegressor using the same data, with the time of event as the label (censored or not) under the exact same settings. The whole fitting process took less than **1 minute**. All common parameters, such as 'n_jobs', were set to be the same, and the only difference was the type of model.  In both cases, 'n_jobs' was set to -1 to utilize parallelization. 

I am struggling to understand why there is such a significant difference in training time between these two models. Although I expect survival analysis to take a bit longer than usual regression, the difference is quite significant. Unfortunately, my entire training dataset has more than 10 million instances, and it seems rather hopeless to apply RandomSurvivalForest at the moment.

I am wondering if there are any suggestions on how I could speed up the training process.
RandomSurvivalForest unusally slow
I attempted to fit a RandomSurvivalForest with 500,000 training instances and 140 features on a machine with 90GB memory. Unfortunately, after **hours** of waiting, the program ran out of memory and crashed. I wasn't able to see any progress even though the parameter 'verbose' was set to 1.

However, I was able to fit a sklearn.RandomForestRegressor using the same data, with the time of event as the label (censored or not) under the exact same settings. The whole fitting process took less than **1 minute**. All common parameters, such as 'n_jobs', were set to be the same, and the only difference was the type of model.  In both cases, 'n_jobs' was set to -1 to utilize parallelization. 

I am struggling to understand why there is such a significant difference in training time between these two models. Although I expect survival analysis to take a bit longer than usual regression, the difference is quite significant. Unfortunately, my entire training dataset has more than 10 million instances, and it seems rather hopeless to apply RandomSurvivalForest at the moment.

I am wondering if there are any suggestions on how I could speed up the training process.


Discussion / hints:
That's not good.

sksurv calls sklearn to grow the forest, it just overrides the split criterion. It would be interesting to check whether increasing the number of samples or the number of features has a higher impact.

I would suspect that the number of samples could be the problem, because computing the split criterion involves sorting by the time of an event.
I can confirm that it is indeed the sorting operation that is responsible for ~60% of the time it takes to fit a tree.

![profiling](https://user-images.githubusercontent.com/74460/224545817-63975367-985e-4f59-8c4e-e9a2c17dc82b.png)

I am also having a very similar experience with RandomSurvivalForest as mentioned by @shil3754 .  I tested it with 100 records and it worked fine. Then I ran it with 500k records on 96 core and 200GB memory machine. Now, It is taking forever to get any output from RandomSurvivalForest even after adding min_sample_splits, max_depth and variours other hyperparameters
@sebp what is your recommendation in view of this? 
> @sebp what is your recommendation in view of this?

I haven't investigated what's happening inside the [tree builder](https://github.com/scikit-learn/scikit-learn/blob/9aaed498795f68e5956ea762fef9c440ca9eb239/sklearn/tree/_splitter.pyx#L561) from sklearn, but my hope is that it would be possible to sort only once and then use the order inside sksurv's split criterion instead of sorting during each call by the tree builder.
having a similar experience with only 33599 samples and 29 features in the training set on a cluster.
FYI, I started looking into this. It should be possible to pre-compute the log-rank statistics for all possible splits in [LogrankCriterion.init](https://github.com/sebp/scikit-survival/blob/d5fcbc9e06f4824f987b5cf550ce0bf82665f28e/sksurv/tree/_criterion.pyx#L129) such that [LogrankCriterion.update](https://github.com/sebp/scikit-survival/blob/d5fcbc9e06f4824f987b5cf550ce0bf82665f28e/sksurv/tree/_criterion.pyx#LL205C14-L205C20) is just a lookup without requiring to sort again.
The [fix-rsf-performance](https://github.com/sebp/scikit-survival/tree/fix-rsf-performance) branch contains improved code to grow trees. For me, it reduces growing a single tree for 500 samples and 25 features from 1330ms to 340ms.

Would be great, if you could give it a try too and let me know what you think.
That's not good.

sksurv calls sklearn to grow the forest, it just overrides the split criterion. It would be interesting to check whether increasing the number of samples or the number of features has a higher impact.

I would suspect that the number of samples could be the problem, because computing the split criterion involves sorting by the time of an event.
I can confirm that it is indeed the sorting operation that is responsible for ~60% of the time it takes to fit a tree.

![profiling](https://user-images.githubusercontent.com/74460/224545817-63975367-985e-4f59-8c4e-e9a2c17dc82b.png)

I am also having a very similar experience with RandomSurvivalForest as mentioned by @shil3754 .  I tested it with 100 records and it worked fine. Then I ran it with 500k records on 96 core and 200GB memory machine. Now, It is taking forever to get any output from RandomSurvivalForest even after adding min_sample_splits, max_depth and variours other hyperparameters
@sebp what is your recommendation in view of this? 
> @sebp what is your recommendation in view of this?

I haven't investigated what's happening inside the [tree builder](https://github.com/scikit-learn/scikit-learn/blob/9aaed498795f68e5956ea762fef9c440ca9eb239/sklearn/tree/_splitter.pyx#L561) from sklearn, but my hope is that it would be possible to sort only once and then use the order inside sksurv's split criterion instead of sorting during each call by the tree builder.
having a similar experience with only 33599 samples and 29 features in the training set on a cluster.
FYI, I started looking into this. It should be possible to pre-compute the log-rank statistics for all possible splits in [LogrankCriterion.init](https://github.com/sebp/scikit-survival/blob/d5fcbc9e06f4824f987b5cf550ce0bf82665f28e/sksurv/tree/_criterion.pyx#L129) such that [LogrankCriterion.update](https://github.com/sebp/scikit-survival/blob/d5fcbc9e06f4824f987b5cf550ce0bf82665f28e/sksurv/tree/_criterion.pyx#LL205C14-L205C20) is just a lookup without requiring to sort again.
The [fix-rsf-performance](https://github.com/sebp/scikit-survival/tree/fix-rsf-performance) branch contains improved code to grow trees. For me, it reduces growing a single tree for 500 samples and 25 features from 1330ms to 340ms.

Would be great, if you could give it a try too and let me know what you think.

Fix patch:
diff --git a/.gitignore b/.gitignore
index 0a11713a..45bf8802 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,6 +8,7 @@ __pycache__/
 # Cython generated files
 _*.c
 _*.cpp
+cython_debug/
 
 # Distribution / packaging
 .Python
diff --git a/doc/user_guide/random-survival-forest.ipynb b/doc/user_guide/random-survival-forest.ipynb
index 0072427a..09b960b3 100644
--- a/doc/user_guide/random-survival-forest.ipynb
+++ b/doc/user_guide/random-survival-forest.ipynb
@@ -376,7 +376,7 @@
     "surv = rsf.predict_survival_function(X_test_sel, return_array=True)\n",
     "\n",
     "for i, s in enumerate(surv):\n",
-    "    plt.step(rsf.event_times_, s, where=\"post\", label=str(i))\n",
+    "    plt.step(rsf.unique_times_, s, where=\"post\", label=str(i))\n",
     "plt.ylabel(\"Survival probability\")\n",
     "plt.xlabel(\"Time in days\")\n",
     "plt.legend()\n",
@@ -412,7 +412,7 @@
     "surv = rsf.predict_cumulative_hazard_function(X_test_sel, return_array=True)\n",
     "\n",
     "for i, s in enumerate(surv):\n",
-    "    plt.step(rsf.event_times_, s, where=\"post\", label=str(i))\n",
+    "    plt.step(rsf.unique_times_, s, where=\"post\", label=str(i))\n",
     "plt.ylabel(\"Cumulative hazard\")\n",
     "plt.xlabel(\"Time in days\")\n",
     "plt.legend()\n",
diff --git a/sksurv/ensemble/boosting.py b/sksurv/ensemble/boosting.py
index e82778ad..6ccc1d2c 100644
--- a/sksurv/ensemble/boosting.py
+++ b/sksurv/ensemble/boosting.py
@@ -161,8 +161,8 @@ class ComponentwiseGradientBoostingSurvivalAnalysis(BaseEnsemble, SurvivalAnalys
         Names of features seen during ``fit``. Defined only when `X`
         has feature names that are all strings.
 
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     References
     ----------
@@ -379,14 +379,14 @@ def predict_cumulative_hazard_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -443,14 +443,14 @@ def predict_survival_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability of
-            survival for each `self.event_times_`, otherwise an array of
+            survival for each `self.unique_times_`, otherwise an array of
             length `n_samples` of :class:`sksurv.functions.StepFunction`
             instances will be returned.
 
@@ -495,7 +495,7 @@ def coef_(self):
         return coef
 
     @property
-    def event_times_(self):
+    def unique_times_(self):
         return self._get_baseline_model().unique_times_
 
     @property
@@ -725,8 +725,8 @@ class GradientBoostingSurvivalAnalysis(BaseGradientBoosting, SurvivalAnalysisMix
         Names of features seen during ``fit``. Defined only when `X`
         has feature names that are all strings.
 
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     References
     ----------
@@ -1272,14 +1272,14 @@ def predict_cumulative_hazard_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -1336,14 +1336,14 @@ def predict_survival_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability of
-            survival for each `self.event_times_`, otherwise an array of
+            survival for each `self.unique_times_`, otherwise an array of
             length `n_samples` of :class:`sksurv.functions.StepFunction`
             instances will be returned.
 
@@ -1379,5 +1379,5 @@ def predict_survival_function(self, X, return_array=False):
         )
 
     @property
-    def event_times_(self):
+    def unique_times_(self):
         return self._get_baseline_model().unique_times_
diff --git a/sksurv/ensemble/forest.py b/sksurv/ensemble/forest.py
index df112f8a..4125fcf5 100644
--- a/sksurv/ensemble/forest.py
+++ b/sksurv/ensemble/forest.py
@@ -19,6 +19,7 @@
 from ..base import SurvivalAnalysisMixin
 from ..metrics import concordance_index_censored
 from ..tree import SurvivalTree
+from ..tree._criterion import get_unique_times
 from ..tree.tree import _array_to_step_function
 from ..util import check_array_survival
 
@@ -91,8 +92,8 @@ def fit(self, X, y, sample_weight=None):
 
         self.n_features_in_ = X.shape[1]
         time = time.astype(np.float64)
-        self.event_times_ = np.unique(time[event])
-        self.n_outputs_ = self.event_times_.shape[0]
+        self.unique_times_, self.is_event_time_ = get_unique_times(time, event)
+        self.n_outputs_ = self.unique_times_.shape[0]
 
         y_numeric = np.empty((X.shape[0], 2), dtype=np.float64)
         y_numeric[:, 0] = time
@@ -139,6 +140,7 @@ def fit(self, X, y, sample_weight=None):
                                           random_state=random_state)
                      for i in range(n_more_estimators)]
 
+            y_tree = y_numeric, self.unique_times_, self.is_event_time_,
             # Parallel loop: we prefer the threading backend as the Cython code
             # for fitting the trees is internally releasing the Python GIL
             # making threading more efficient than multiprocessing in
@@ -148,7 +150,7 @@ def fit(self, X, y, sample_weight=None):
             trees = Parallel(n_jobs=self.n_jobs, verbose=self.verbose,
                              prefer='threads')(
                 delayed(_parallel_build_trees)(
-                    t, self.bootstrap, X, (y_numeric, self.event_times_), sample_weight, i, len(trees),
+                    t, self.bootstrap, X, y_tree, sample_weight, i, len(trees),
                     verbose=self.verbose,
                     n_samples_bootstrap=n_samples_bootstrap)
                 for i, t in enumerate(trees))
@@ -258,13 +260,13 @@ def predict_cumulative_hazard_function(self, X, return_array=False):
         arr = self._predict("predict_cumulative_hazard_function", X)
         if return_array:
             return arr
-        return _array_to_step_function(self.event_times_, arr)
+        return _array_to_step_function(self.unique_times_, arr)
 
     def predict_survival_function(self, X, return_array=False):
         arr = self._predict("predict_survival_function", X)
         if return_array:
             return arr
-        return _array_to_step_function(self.event_times_, arr)
+        return _array_to_step_function(self.unique_times_, arr)
 
 
 class RandomSurvivalForest(_BaseSurvivalForest):
@@ -380,8 +382,8 @@ class RandomSurvivalForest(_BaseSurvivalForest):
     estimators_ : list of SurvivalTree instances
         The collection of fitted sub-estimators.
 
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     n_features_in_ : int
         Number of features seen during ``fit``.
@@ -497,14 +499,14 @@ def predict_cumulative_hazard_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -555,14 +557,14 @@ def predict_survival_function(self, X, return_array=False):
 
         return_array : boolean
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`
             will be returned.
 
@@ -713,8 +715,8 @@ class ExtraSurvivalTrees(_BaseSurvivalForest):
     estimators_ : list of SurvivalTree instances
         The collection of fitted sub-estimators.
 
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     n_features_in_ : int
         The number of features when ``fit`` is performed.
@@ -799,14 +801,14 @@ def predict_cumulative_hazard_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -857,14 +859,14 @@ def predict_survival_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability of
-            survival for each `self.event_times_`, otherwise an array of
+            survival for each `self.unique_times_`, otherwise an array of
             length `n_samples` of :class:`sksurv.functions.StepFunction`
             instances will be returned.
 
diff --git a/sksurv/linear_model/coxnet.py b/sksurv/linear_model/coxnet.py
index 0d086493..0b4de2f2 100644
--- a/sksurv/linear_model/coxnet.py
+++ b/sksurv/linear_model/coxnet.py
@@ -124,8 +124,8 @@ class CoxnetSurvivalAnalysis(BaseEstimator, SurvivalAnalysisMixin):
         Names of features seen during ``fit``. Defined only when `X`
         has feature names that are all strings.
 
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     References
     ----------
@@ -392,14 +392,14 @@ def predict_cumulative_hazard_function(self, X, alpha=None, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -469,14 +469,14 @@ def predict_survival_function(self, X, alpha=None, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability of
-            survival for each `self.event_times_`, otherwise an array of
+            survival for each `self.unique_times_`, otherwise an array of
             length `n_samples` of :class:`sksurv.functions.StepFunction`
             instances will be returned.
 
@@ -522,5 +522,5 @@ def predict_survival_function(self, X, alpha=None, return_array=False):
         )
 
     @property
-    def event_times_(self):
+    def unique_times_(self):
         return self._get_baseline_model(None).unique_times_
diff --git a/sksurv/linear_model/coxph.py b/sksurv/linear_model/coxph.py
index c946bb9f..f6d45826 100644
--- a/sksurv/linear_model/coxph.py
+++ b/sksurv/linear_model/coxph.py
@@ -351,8 +351,8 @@ class CoxPHSurvivalAnalysis(BaseEstimator, SurvivalAnalysisMixin):
         Names of features seen during ``fit``. Defined only when `X`
         has feature names that are all strings.
 
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     See also
     --------
@@ -395,7 +395,7 @@ def baseline_survival_(self):
         return self._baseline_model.baseline_survival_
 
     @property
-    def event_times_(self):
+    def unique_times_(self):
         return self._baseline_model.unique_times_
 
     def fit(self, X, y):
@@ -522,14 +522,14 @@ def predict_cumulative_hazard_function(self, X, return_array=False):
 
         return_array : boolean
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -583,14 +583,14 @@ def predict_survival_function(self, X, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability of
-            survival for each `self.event_times_`, otherwise an array of
+            survival for each `self.unique_times_`, otherwise an array of
             length `n_samples` of :class:`sksurv.functions.StepFunction`
             instances will be returned.
 
diff --git a/sksurv/tree/_criterion.pyx b/sksurv/tree/_criterion.pyx
index b5c3fabc..f756fc44 100644
--- a/sksurv/tree/_criterion.pyx
+++ b/sksurv/tree/_criterion.pyx
@@ -2,8 +2,8 @@
 # cython: boundscheck=False
 # cython: wraparound=False
 
-from libc.math cimport INFINITY, fabs, sqrt
-from libc.stdlib cimport calloc, free, malloc, qsort, realloc
+from libc.math cimport INFINITY, NAN, fabs, sqrt
+from libc.stdlib cimport free, malloc
 from libc.string cimport memset
 
 import numpy as np
@@ -15,89 +15,131 @@ cnp.import_array()
 from sklearn.tree._criterion cimport Criterion
 from sklearn.tree._tree cimport DOUBLE_t, SIZE_t
 
-ctypedef struct Timepoint:
 
-    SIZE_t index
-    double time
-    double event
-
-
-ctypedef struct Riskset:
-
-    double n_events
-    SIZE_t n_at_risk
+cpdef get_unique_times(cnp.ndarray[DOUBLE_t, ndim=1] time, cnp.ndarray[cnp.npy_bool, ndim=1] event):
+    cdef:
+        SIZE_t[:] order = cnp.PyArray_ArgSort(time, 0, cnp.NPY_MERGESORT)
+        DOUBLE_t value
+        DOUBLE_t last_value = NAN
+        SIZE_t i
+        SIZE_t idx
+        list unique_values = []
+        list has_event = []
+
+    for i in range(time.shape[0]):
+        idx = order[i]
+        value = time[idx]
+        if value != last_value:
+            unique_values.append(value)
+            has_event.append(event[idx])
+            last_value = value
+        if event[idx]:
+            has_event[len(has_event) - 1] = True
+
+    return np.asarray(unique_values), np.asarray(has_event, dtype=np.bool_)
+
+
+cdef class RisksetCounter:
+    cdef:
+        const DOUBLE_t[:] unique_times
+        cnp.npy_int64 * n_events
+        cnp.npy_int64 * n_at_risk
+        const DOUBLE_t[:, ::1] data
+        size_t nbytes
+
+    def __cinit__(self, const DOUBLE_t[:] unique_times):
+        cdef SIZE_t n_unique_times = unique_times.shape[0]
+        self.nbytes = n_unique_times * sizeof(cnp.npy_int64)
+        self.n_events = <cnp.npy_int64 *> malloc(self.nbytes)
+        self.n_at_risk = <cnp.npy_int64 *> malloc(self.nbytes)
+        self.unique_times = unique_times
 
+    def __dealloc__(self):
+        """Destructor."""
+        free(self.n_events)
+        free(self.n_at_risk)
+
+    cdef void reset(self) nogil:
+        memset(self.n_events, 0, self.nbytes)
+        memset(self.n_at_risk, 0, self.nbytes)
+
+    cdef void set_data(self, const DOUBLE_t[:, ::1] data) nogil:
+        self.data = data
+
+    cdef void update(self, const SIZE_t* samples, SIZE_t start, SIZE_t end) nogil:
+        cdef:
+            SIZE_t i
+            SIZE_t idx
+            SIZE_t ti
+            DOUBLE_t time
+            DOUBLE_t event
+            const DOUBLE_t[:] unique_times = self.unique_times
+            SIZE_t n_times = unique_times.shape[0]
+            const DOUBLE_t[:, ::1] y = self.data
 
-cdef int compare_timepoint_desc(const void* a, const void* b) nogil:
-    """Comparison function for sort by time (descending order)."""
-    cdef double ta = (<Timepoint *> a).time
-    cdef double tb = (<Timepoint *> b).time
+        self.reset()
 
-    if ta > tb:
-        return -1
-    if ta < tb:
-        return 1
-    return 0
+        for i in range(start, end):
+            idx = samples[i]
+            time, event = y[idx, 0], y[idx, 1]
 
+            # i-th sample is in all risk sets with time <= i-th time
+            ti = 0
+            while ti < n_times and unique_times[ti] < time:
+                self.n_at_risk[ti] += 1
+                ti += 1
 
-cdef inline Riskset* set_risk_set(Riskset * riskset, SIZE_t n_events, SIZE_t n_at_risk) nogil:
-    riskset.n_events = <double> n_events
-    riskset.n_at_risk = n_at_risk
-    return riskset
+            if ti < n_times:  # unique_times[ti] == time
+                self.n_at_risk[ti] += 1
+                if event != 0.0:
+                    self.n_events[ti] += 1
 
+    cdef inline void at(self, SIZE_t index, DOUBLE_t * at_risk, DOUBLE_t * events) nogil:
+        if at_risk != NULL:
+            at_risk[0] = <DOUBLE_t> self.n_at_risk[index]
+        if events != NULL:
+            events[0] = <DOUBLE_t> self.n_events[index]
 
-cdef inline Timepoint* set_time_point(Timepoint * tp_ptr, SIZE_t i, const DOUBLE_t[:, ::1] y) nogil:
-    tp_ptr.index = i
-    tp_ptr.time = y[i, 0]
-    tp_ptr.event = y[i, 1]
-    return tp_ptr
 
+cdef int argbinsearch(const DOUBLE_t[:] arr, DOUBLE_t key_val, SIZE_t * ret) nogil except -1:
+    cdef:
+        SIZE_t arr_len = arr.shape[0]
+        SIZE_t min_idx = 0
+        SIZE_t max_idx = arr_len
+        SIZE_t mid_idx
+        DOUBLE_t mid_val
 
-cdef void compute_risksets(const Timepoint * time_arr, SIZE_t n_samples,
-                           const DOUBLE_t[::1] event_times,
-                           Riskset * riskset) nogil:
-    cdef DOUBLE_t time_idx
-    cdef DOUBLE_t tp
-    cdef SIZE_t total_events
-    cdef SIZE_t n_event_times = event_times.shape[0]
-    cdef SIZE_t idx = 0  # index over time_arr
-    cdef SIZE_t idx_tp = n_event_times - 1  # index over event_times and riskset
+    while min_idx < max_idx:
+        mid_idx = min_idx + ((max_idx - min_idx) >> 1)
 
-    while idx_tp >= 0 and idx < n_samples:
-        time_idx = time_arr[idx].time
-        tp = event_times[idx_tp]
-        if time_idx < tp:
-            set_risk_set(&riskset[idx_tp], 0, idx)
-            idx_tp -= 1  # move to next smaller event time point
-            continue
-        if time_idx > tp:
-            idx += 1  # move to to next smaller sample time point
-            continue
+        if mid_idx < 0 or mid_idx >= arr_len:
+            return -1
 
-        total_events = 0
-        while idx < n_samples and tp == time_arr[idx].time:
-            if time_arr[idx].event != 0.0:
-                total_events += 1
-            idx += 1
+        mid_val = arr[mid_idx]
+        if mid_val < key_val:
+            min_idx = mid_idx + 1
+        else:
+            max_idx = mid_idx
 
-        set_risk_set(&riskset[idx_tp], total_events, idx)
-        idx_tp -= 1  # move to next smaller event time point
+    ret[0] = min_idx
 
-    # for remaining smaller time points, everyone is at risk
-    while idx_tp >= 0:
-        set_risk_set(&riskset[idx_tp], 0, n_samples)
-        idx_tp -= 1
+    return 0
 
 
 cdef class LogrankCriterion(Criterion):
 
-    # unique time points sorted in ascending order
-    cdef const DOUBLE_t[::1] event_times
-    cdef SIZE_t n_event_times
-    cdef Riskset * riskset_total
-    cdef Riskset * riskset_left
-
-    def __cinit__(self, SIZE_t n_outputs, SIZE_t n_samples, const DOUBLE_t[::1] event_times):
+    cdef:
+        # unique time points sorted in ascending order
+        const DOUBLE_t[::1] unique_times
+        SIZE_t n_unique_times
+        size_t nbytes
+        RisksetCounter riskset_total
+        cnp.npy_int64 * delta_n_at_risk_left
+        cnp.npy_int64 * n_events_left
+        SIZE_t * samples_time_idx
+        SIZE_t n_samples_left
+
+    def __cinit__(self, SIZE_t n_outputs, SIZE_t n_samples, const DOUBLE_t[::1] unique_times):
         # Default values
         self.samples = NULL
         self.start = 0
@@ -106,25 +148,27 @@ cdef class LogrankCriterion(Criterion):
 
         self.n_outputs = n_outputs
         self.n_samples = n_samples
-        self.event_times = event_times
-        self.n_event_times = event_times.shape[0]
+        self.unique_times = unique_times
+        self.n_unique_times = unique_times.shape[0]
+        self.nbytes = self.n_unique_times * sizeof(cnp.npy_int64)
         self.n_node_samples = 0
         self.weighted_n_node_samples = 0.0
         self.weighted_n_left = 0.0
         self.weighted_n_right = 0.0
 
-        self.riskset_total = NULL
-        self.riskset_left = NULL
+        self.riskset_total = RisksetCounter(unique_times)
+        self.delta_n_at_risk_left = <cnp.npy_int64 *> malloc(self.nbytes)
+        self.n_events_left = <cnp.npy_int64 *> malloc(self.nbytes)
+        self.samples_time_idx = <SIZE_t *> malloc(n_samples * sizeof(SIZE_t))
 
     def __dealloc__(self):
         """Destructor."""
-        cdef SIZE_t p
-        if self.riskset_total is not NULL:
-            free(self.riskset_total)
-            free(self.riskset_left)
+        free(self.delta_n_at_risk_left)
+        free(self.n_events_left)
+        free(self.samples_time_idx)
 
     def __reduce__(self):
-        return (type(self), (self.n_outputs, self.n_samples, self.event_times), self.__getstate__())
+        return (type(self), (self.n_outputs, self.n_samples, self.unique_times), self.__getstate__())
 
     cdef int init(self, const DOUBLE_t[:, ::1] y, const DOUBLE_t[:] sample_weight,
                   double weighted_n_samples, SIZE_t* samples, SIZE_t start,
@@ -141,52 +185,32 @@ cdef class LogrankCriterion(Criterion):
         self.weighted_n_samples = weighted_n_samples
         self.weighted_n_node_samples = 0.
 
-        cdef SIZE_t i
-        cdef SIZE_t p
-        cdef SIZE_t k
-        cdef DOUBLE_t w = 1.0
-        cdef Timepoint * time_arr
+        cdef:
+            SIZE_t i
+            SIZE_t idx
+            DOUBLE_t time
+            DOUBLE_t w = 1.0
+            const DOUBLE_t[::1] unique_times = self.unique_times
 
-        time_arr = < Timepoint * > malloc(self.n_node_samples * sizeof(Timepoint))
-        if time_arr is NULL:
-            raise MemoryError()
+        self.riskset_total.set_data(y)
+        self.riskset_total.update(samples, start, end)
 
-        for k, p in enumerate(range(start, end)):
-            i = samples[p]
-            set_time_point(&time_arr[k], i, self.y)
+        for i in range(start, end):
+            idx = samples[i]
+            time = y[idx, 0]
+            argbinsearch(unique_times, time, &self.samples_time_idx[idx])
 
             if sample_weight is not None:
-                w = sample_weight[i]
+                w = sample_weight[idx]
 
             self.weighted_n_node_samples += w
 
-        qsort(time_arr, self.n_node_samples, sizeof(Timepoint), compare_timepoint_desc)
-
-        if self.riskset_total is not NULL:
-            free(self.riskset_total)
-            free(self.riskset_left)
-
-        self.riskset_total = < Riskset * > malloc(self.n_event_times * sizeof(Riskset))
-        if self.riskset_total is NULL:
-            raise MemoryError()
-
-        compute_risksets(time_arr, self.n_node_samples, self.event_times, self.riskset_total)
-
-        free(time_arr)
-
-        self.riskset_left = < Riskset * > malloc(self.n_event_times * sizeof(Riskset))
-        if self.riskset_left is NULL:
-            raise MemoryError()
-
         # Reset to pos=start
         self.reset()
         return 0
 
     cdef int reset(self) nogil except -1:
         """Reset the criterion at pos=start."""
-        # cdef SIZE_t n_bytes = self.n_samples * sizeof(double)
-        # memset(self.time_points, 0, n_bytes)
-
         self.weighted_n_left = 0.0
         self.weighted_n_right = self.weighted_n_node_samples
         self.pos = self.start
@@ -194,9 +218,6 @@ cdef class LogrankCriterion(Criterion):
 
     cdef int reverse_reset(self) nogil except -1:
         """Reset the criterion at pos=end."""
-        # cdef SIZE_t n_bytes = self.n_samples * sizeof(double)
-        # memset(self.time_points, 0, n_bytes)
-
         self.weighted_n_right = 0.0
         self.weighted_n_left = self.weighted_n_node_samples
         self.pos = self.end
@@ -204,41 +225,38 @@ cdef class LogrankCriterion(Criterion):
 
     cdef int update(self, SIZE_t new_pos) nogil except -1:
         """Updated statistics by moving samples[pos:new_pos] to the left."""
-        cdef const DOUBLE_t[:] sample_weight = self.sample_weight
-        cdef const SIZE_t* samples = self.samples
-
-        cdef SIZE_t pos = self.start  # always start from the beginning
-        cdef SIZE_t n_samples_left = new_pos - pos
-        cdef SIZE_t i
-        cdef SIZE_t idx_left
-        cdef SIZE_t k
-        cdef DOUBLE_t w = 1.0
-        cdef Timepoint * time_arr
-
-        time_arr = < Timepoint * > malloc(n_samples_left * sizeof(Timepoint))
-        if time_arr is NULL:
-            raise MemoryError()
+        cdef:
+            const DOUBLE_t[:] sample_weight = self.sample_weight
+            const SIZE_t* samples = self.samples
+            const DOUBLE_t[:, ::1] y = self.y
+
+            SIZE_t pos = self.start  # always start from the beginning
+            SIZE_t i
+            SIZE_t idx
+            DOUBLE_t event
+            SIZE_t time_idx
+            DOUBLE_t w = 1.0
+
+        self.n_samples_left = new_pos - pos
+        memset(self.delta_n_at_risk_left, 0, self.nbytes)
+        memset(self.n_events_left, 0, self.nbytes)
 
         # Update statistics up to new_pos
         self.weighted_n_left = 0.0
-        for k, idx_left in enumerate(range(pos, new_pos)):
-            i = samples[idx_left]
-            set_time_point(&time_arr[k], i, self.y)
+        for i in range(pos, new_pos):
+            idx = samples[i]
+            event = y[idx, 1]
+            time_idx = self.samples_time_idx[idx]
+
+            self.delta_n_at_risk_left[time_idx] += 1
+            if event != 0.0:
+                self.n_events_left[time_idx] += 1
 
             if sample_weight is not None:
-                w = sample_weight[i]
+                w = sample_weight[idx]
 
             self.weighted_n_left += w
 
-        qsort(time_arr, n_samples_left, sizeof(Timepoint), compare_timepoint_desc)
-
-        memset(self.riskset_left, 0, self.n_event_times * sizeof(Riskset))
-
-        # use same time points as in riskset_total
-        compute_risksets(time_arr, n_samples_left, self.event_times, self.riskset_left)
-
-        free(time_arr)
-
         self.weighted_n_right = (self.weighted_n_node_samples -
                                  self.weighted_n_left)
 
@@ -254,32 +272,31 @@ cdef class LogrankCriterion(Criterion):
     cdef double proxy_impurity_improvement(self) nogil:
         """Compute a proxy of the impurity reduction"""
 
-        # cdef double impurity_left
-        # cdef double impurity_right
-        #
-        # self.children_impurity(&impurity_left, &impurity_right)
-
-        cdef SIZE_t i
-        cdef double at_risk
-        cdef double total_at_risk
-        cdef double v
-        cdef Riskset * rs_total
-        cdef Riskset * rs_left
-        cdef double denom = 0.0
-        cdef double numer = 0.0
-
-        for i in range(self.n_event_times):
-            rs_total = &self.riskset_total[i]
-            rs_left = &self.riskset_left[i]
-            total_at_risk = <double> rs_total.n_at_risk
+        cdef:
+            SIZE_t i
+            DOUBLE_t at_risk = <DOUBLE_t> self.n_samples_left
+            DOUBLE_t events
+            DOUBLE_t total_at_risk
+            DOUBLE_t total_events
+            DOUBLE_t ratio
+            DOUBLE_t v
+            DOUBLE_t denom = 0.0
+            DOUBLE_t numer = 0.0
+
+        for i in range(self.n_unique_times):
+            events = <DOUBLE_t> self.n_events_left[i]
+            self.riskset_total.at(i, &total_at_risk, &total_events)
 
             if total_at_risk == 0:
                 break  # we reached the end
-            at_risk = rs_left.n_at_risk / total_at_risk
-            numer += rs_left.n_events - rs_total.n_events * at_risk
+            ratio = at_risk / total_at_risk
+            numer += events - total_events * ratio
             if total_at_risk > 1.0:
-                v = (total_at_risk - rs_total.n_events) / (total_at_risk - 1.0) * rs_total.n_events
-                denom += at_risk * (1.0 - at_risk) * v
+                v = (total_at_risk - total_events) / (total_at_risk - 1.0) * total_events
+                denom += ratio * (1.0 - ratio) * v
+
+            # Update number of samples at risk for next bigger timepoint
+            at_risk -= <DOUBLE_t> self.delta_n_at_risk_left[i]
 
         if denom != 0.0:
             # absolute value is the measure of node separation
@@ -305,23 +322,25 @@ cdef class LogrankCriterion(Criterion):
     cdef void node_value(self, double* dest) nogil:
         """Compute the node value of samples[start:end] into dest."""
         # Estimate cumulative hazard function
-        cdef SIZE_t k
-        cdef SIZE_t j
-        cdef double ratio
-        cdef Riskset * rs
-
-        rs = &self.riskset_total[0]
-        ratio = rs.n_events / (<double> rs.n_at_risk)
+        cdef:
+            SIZE_t i
+            SIZE_t j
+            DOUBLE_t ratio
+            DOUBLE_t n_events
+            DOUBLE_t n_at_risk
+
+        self.riskset_total.at(0, &n_at_risk, &n_events)
+        ratio = n_events / n_at_risk
         dest[0] = ratio  # Nelson-Aalen estimator
         dest[1] = 1.0 - ratio  # Kaplan-Meier estimator
 
         j = 2
-        for k in range(1, self.n_event_times):
-            rs = &self.riskset_total[k]
+        for i in range(1, self.n_unique_times):
+            self.riskset_total.at(i, &n_at_risk, &n_events)
             dest[j] = dest[j - 2]
             dest[j + 1] = dest[j - 1]
-            if rs.n_at_risk != 0:
-                ratio = rs.n_events / (<double> rs.n_at_risk)
+            if n_at_risk != 0:
+                ratio = n_events / n_at_risk
                 dest[j] += ratio
                 dest[j + 1] *= 1.0 - ratio
             j += 2
diff --git a/sksurv/tree/tree.py b/sksurv/tree/tree.py
index 991667c1..2027e8fd 100644
--- a/sksurv/tree/tree.py
+++ b/sksurv/tree/tree.py
@@ -14,7 +14,7 @@
 from ..base import SurvivalAnalysisMixin
 from ..functions import StepFunction
 from ..util import check_array_survival
-from ._criterion import LogrankCriterion
+from ._criterion import LogrankCriterion, get_unique_times
 
 __all__ = ["SurvivalTree"]
 
@@ -109,8 +109,8 @@ class SurvivalTree(BaseEstimator, SurvivalAnalysisMixin):
 
     Attributes
     ----------
-    event_times_ : array of shape = (n_event_times,)
-        Unique time points where events occurred.
+    unique_times_ : array of shape = (n_unique_times,)
+        Unique time points.
 
     max_features_ : int,
         The inferred value of max_features.
@@ -211,7 +211,7 @@ def fit(self, X, y, sample_weight=None, check_input=True):
             X = self._validate_data(X, ensure_min_samples=2, accept_sparse="csc")
             event, time = check_array_survival(X, y)
             time = time.astype(np.float64)
-            self.event_times_ = np.unique(time[event])
+            self.unique_times_, self.is_event_time_ = get_unique_times(time, event)
             if issparse(X):
                 X.sort_indices()
 
@@ -219,18 +219,18 @@ def fit(self, X, y, sample_weight=None, check_input=True):
             y_numeric[:, 0] = time
             y_numeric[:, 1] = event.astype(np.float64)
         else:
-            y_numeric, self.event_times_ = y
+            y_numeric, self.unique_times_, self.is_event_time_ = y
 
         n_samples, self.n_features_in_ = X.shape
         params = self._check_params(n_samples)
 
-        self.n_outputs_ = self.event_times_.shape[0]
+        self.n_outputs_ = self.unique_times_.shape[0]
         # one "class" for CHF, one for survival function
         self.n_classes_ = np.ones(self.n_outputs_, dtype=np.intp) * 2
 
         # Build tree
         self.criterion = "logrank"
-        criterion = LogrankCriterion(self.n_outputs_, n_samples, self.event_times_)
+        criterion = LogrankCriterion(self.n_outputs_, n_samples, self.unique_times_)
 
         SPLITTERS = SPARSE_SPLITTERS if issparse(X) else DENSE_SPLITTERS
 
@@ -370,7 +370,7 @@ def predict(self, X, check_input=True):
             Predicted risk scores.
         """
         chf = self.predict_cumulative_hazard_function(X, check_input, return_array=True)
-        return chf.sum(1)
+        return chf[:, self.is_event_time_].sum(1)
 
     def predict_cumulative_hazard_function(self, X, check_input=True, return_array=False):
         """Predict cumulative hazard function.
@@ -392,14 +392,14 @@ def predict_cumulative_hazard_function(self, X, check_input=True, return_array=F
 
         return_array : boolean, default: False
             If set, return an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of
+            for each `self.unique_times_`, otherwise an array of
             :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         cum_hazard : ndarray
             If `return_array` is set, an array with the cumulative hazard rate
-            for each `self.event_times_`, otherwise an array of length `n_samples`
+            for each `self.unique_times_`, otherwise an array of length `n_samples`
             of :class:`sksurv.functions.StepFunction` instances will be returned.
 
         Examples
@@ -436,7 +436,7 @@ def predict_cumulative_hazard_function(self, X, check_input=True, return_array=F
         arr = pred[..., 0]
         if return_array:
             return arr
-        return _array_to_step_function(self.event_times_, arr)
+        return _array_to_step_function(self.unique_times_, arr)
 
     def predict_survival_function(self, X, check_input=True, return_array=False):
         """Predict survival function.
@@ -458,14 +458,14 @@ def predict_survival_function(self, X, check_input=True, return_array=False):
 
         return_array : boolean, default: False
             If set, return an array with the probability
-            of survival for each `self.event_times_`,
+            of survival for each `self.unique_times_`,
             otherwise an array of :class:`sksurv.functions.StepFunction`.
 
         Returns
         -------
         survival : ndarray
             If `return_array` is set, an array with the probability of
-            survival for each `self.event_times_`, otherwise an array of
+            survival for each `self.unique_times_`, otherwise an array of
             length `n_samples` of :class:`sksurv.functions.StepFunction`
             instances will be returned.
 
@@ -503,7 +503,7 @@ def predict_survival_function(self, X, check_input=True, return_array=False):
         arr = pred[..., 1]
         if return_array:
             return arr
-        return _array_to_step_function(self.event_times_, arr)
+        return _array_to_step_function(self.unique_times_, arr)
 
     def apply(self, X, check_input=True):
         """Return the index of the leaf that each sample is predicted as.


Test patch:
diff --git a/tests/test_forest.py b/tests/test_forest.py
index 5b53c798..5fa136c2 100644
--- a/tests/test_forest.py
+++ b/tests/test_forest.py
@@ -52,7 +52,7 @@ def test_fit_int_time(make_whas500, forest_cls):
     assert len(forest_f.estimators_) == len(forest_i.estimators_)
     assert forest_f.n_features_in_ == forest_i.n_features_in_
     assert forest_f.oob_score_ == forest_i.oob_score_
-    assert_array_almost_equal(forest_f.event_times_, forest_i.event_times_)
+    assert_array_almost_equal(forest_f.unique_times_, forest_i.unique_times_)
 
     pred_f = forest_f.predict(whas500.x[:50])
     pred_i = forest_i.predict(whas500.x[:50])
@@ -70,7 +70,7 @@ def test_fit_predict_chf(make_whas500, forest_cls):
     assert len(forest.estimators_) == 10
 
     chf = forest.predict_cumulative_hazard_function(whas500.x, return_array=True)
-    assert chf.shape == (500, forest.event_times_.shape[0])
+    assert chf.shape == (500, forest.unique_times_.shape[0])
 
     assert np.isfinite(chf).all()
     assert np.all(chf >= 0.0)
@@ -93,7 +93,7 @@ def test_fit_predict_surv(make_whas500, forest_cls):
     assert len(forest.estimators_) == 10
 
     surv = forest.predict_survival_function(whas500.x, return_array=True)
-    assert surv.shape == (500, forest.event_times_.shape[0])
+    assert surv.shape == (500, forest.unique_times_.shape[0])
 
     assert np.isfinite(surv).all()
     assert np.all(surv >= 0.0)
@@ -143,7 +143,7 @@ def test_predict_step_function(make_whas500, forest_cls, func):
     assert ret_array.shape[0] == fn_array.shape[0]
 
     for fn, arr in zip(fn_array, ret_array):
-        assert_array_almost_equal(fn.x, forest.event_times_)
+        assert_array_almost_equal(fn.x, forest.unique_times_)
         assert_array_almost_equal(fn.y, arr)
 
 
diff --git a/tests/test_survival_function.py b/tests/test_survival_function.py
index 2d9c58de..5449700a 100644
--- a/tests/test_survival_function.py
+++ b/tests/test_survival_function.py
@@ -26,7 +26,7 @@ def test_survival_functions(estimator, make_whas500):
     fns_cls = estimator.predict_survival_function(data.x[:150])
     fns_arr = estimator.predict_survival_function(data.x[:150], return_array=True)
 
-    times = estimator.event_times_
+    times = estimator.unique_times_
     arr = np.row_stack([fn(times) for fn in fns_cls])
 
     assert_array_almost_equal(arr, fns_arr)
diff --git a/tests/test_tree.py b/tests/test_tree.py
index 5ff35957..42b34f65 100644
--- a/tests/test_tree.py
+++ b/tests/test_tree.py
@@ -68,6 +68,7 @@ def assert_curve_almost_equal(x, y):
     jumps_x = np.diff(x) != 0
     jumps_y = np.diff(y) != 0
 
+    pytest.approx(y[0], x[0])
     assert_array_almost_equal(x[:1], y[:1])
     assert_array_almost_equal(x[1:][jumps_x], y[1:][jumps_y])
 
@@ -165,13 +166,14 @@ def test_tree_one_split(veterans):
         1, 2, 3, 4, 7, 8, 10, 11, 12, 13, 15, 16, 18, 19, 20,
         21, 22, 24, 25, 27, 29, 30, 31, 33, 35, 36, 42, 43, 44,
         45, 48, 49, 51, 52, 53, 54, 56, 59, 61, 63, 72, 73, 80,
-        82, 84, 87, 90, 92, 95, 99, 100, 103, 105, 110, 111, 112,
-        117, 118, 122, 126, 132, 133, 139, 140, 143, 144, 151, 153,
-        156, 162, 164, 177, 186, 200, 201, 216, 228, 231, 242, 250,
+        82, 83, 84, 87, 90, 92, 95, 97, 99, 100, 103, 105, 110, 111, 112,
+        117, 118, 122, 123, 126, 132, 133, 139, 140, 143, 144, 151, 153,
+        156, 162, 164, 177, 182, 186, 200, 201, 216, 228, 231, 242, 250,
         260, 278, 283, 287, 314, 340, 357, 378, 384, 389, 392, 411,
         467, 553, 587, 991, 999,
     ], dtype=float)
-    assert_array_equal(tree.event_times_, expected_time)
+    assert_array_equal(tree.unique_times_, expected_time)
+    assert_array_equal(tree.unique_times_[~tree.is_event_time_], np.array([83, 97, 123, 182], dtype=float))
 
     threshold = stats.loc[0, "threshold"]
     m = X[:, 0] <= threshold
@@ -349,7 +351,7 @@ def test_fit_int_time(breast_cancer):
         random_state=6,
     ).fit(X, y_int)
 
-    assert_array_almost_equal(tree_f.event_times_, tree_i.event_times_)
+    assert_array_almost_equal(tree_f.unique_times_, tree_i.unique_times_)
     assert_array_equal(tree_f.tree_.feature, tree_i.tree_.feature)
     assert_array_equal(tree_f.tree_.n_node_samples, tree_i.tree_.n_node_samples)
     assert_array_almost_equal(tree_f.tree_.threshold, tree_i.tree_.threshold)
@@ -377,7 +379,7 @@ def test_predict_step_function(breast_cancer, func):
     assert ret_array.shape[0] == fn_array.shape[0]
 
     for fn, arr in zip(fn_array, ret_array):
-        assert_array_almost_equal(fn.x, tree.event_times_)
+        assert_array_almost_equal(fn.x, tree.unique_times_)
         assert_array_almost_equal(fn.y, arr)