Skip to content

Commit ee2ab39

Browse files
authored
MNT: use builtin nsmallest (#633)
* MNT: use builtin nsmallest Signed-off-by: DerWeh <andreas.weh@web.de> * MNT: use stable search order to keep exact behavior Signed-off-by: DerWeh <andreas.weh@web.de> --------- Signed-off-by: DerWeh <andreas.weh@web.de>
1 parent 0917b29 commit ee2ab39

File tree

1 file changed

+5
-10
lines changed
  • python/interpret-core/interpret/glassbox/_ebm

1 file changed

+5
-10
lines changed

python/interpret-core/interpret/glassbox/_ebm/_ebm.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1343,16 +1343,11 @@ def fit(self, X, y, sample_weight=None, bags=None, init_score=None):
13431343
(rank - old_mean) / (n + 1)
13441344
)
13451345

1346-
final_ranks = []
1347-
total_interactions = 0
1348-
for indices in pair_ranks:
1349-
heapq.heappush(final_ranks, (pair_ranks[indices], indices))
1350-
total_interactions += 1
1351-
1352-
n_interactions = min(interactions, total_interactions)
1353-
boost_groups = [
1354-
heapq.heappop(final_ranks)[1] for _ in range(n_interactions)
1355-
]
1346+
boost_groups = heapq.nsmallest(
1347+
interactions,
1348+
pair_ranks,
1349+
key=lambda indices: (pair_ranks[indices], indices),
1350+
)
13561351
else:
13571352
# Check and remove duplicate interaction terms
13581353
uniquifier = set()

0 commit comments

Comments
 (0)