MCPcopy Create free account
hub / github.com/apple/axlearn / _find_first

Method _find_first

axlearn/common/segment_tree.py:113–142  ·  view source on GitHub ↗

Recursive implementation of find_first.

(
        self, node: int, start: int, end: int, predicate: Callable[[T], bool], search_start: int
    )

Source from the content-addressed store, hash-verified

111 return self._find_first(1, 0, self.tree_size - 1, predicate, start_index)
112
113 def _find_first(
114 self, node: int, start: int, end: int, predicate: Callable[[T], bool], search_start: int
115 ) -> int:
116 """Recursive implementation of find_first."""
117 if end < search_start:
118 return -1
119
120 # Leaf node
121 if start == end:
122 if start >= search_start and predicate(self.tree[node]):
123 return start
124 return -1
125
126 mid = (start + end) // 2
127 left_child = 2 * node
128 right_child = 2 * node + 1
129
130 # Search left subtree first (for smallest index)
131 if mid >= search_start and predicate(self.tree[left_child]):
132 result = self._find_first(left_child, start, mid, predicate, search_start)
133 if result != -1:
134 return result
135
136 # Search right subtree
137 if predicate(self.tree[right_child]):
138 return self._find_first(
139 right_child, mid + 1, end, predicate, max(search_start, mid + 1)
140 )
141
142 return -1
143
144 def get_value(self, index: int) -> T:
145 """Get value at a specific index."""

Callers 1

find_firstMethod · 0.95

Calls 1

predicateFunction · 0.85

Tested by

no test coverage detected