added AdjacencyList.leaves #99

jacanchaplais · Feb 22, 2023 · 27ecdf9 · 27ecdf9
1 parent 0d60ca4
commit 27ecdf9
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/graphicle/data.py b/graphicle/data.py
@@ -1852,6 +1852,16 @@ def _edge_relabel(self) -> base.IntVector:
         _, inv = np.unique(self._data, return_inverse=True)
         return inv.reshape(-1, 2)
 
+    @property
+    def leaves(self) -> MaskArray:
+        """A mask to select the leaves of the DAG."""
+        sparse_copy = self._sparse.copy()
+        sparse_copy.data[...] = True
+        out_degree = sparse_copy.sum(axis=1)
+        zero_idxs = np.flatnonzero(out_degree == 0)
+        leaf_mask = np.in1d(sparse_copy.col, zero_idxs)
+        return MaskArray(leaf_mask)
+
     @property
     def data(self) -> base.VoidVector:
         return self._data.view(self.dtype).reshape(-1)