I'm trying to split a sorted integer list into two lists. The first list would have all ints under n
and the second all ints over n
. Note that n
does not have to be in the original list.
I can easily do this with:
under = []
over = []
for x in sorted_list:
if x < n:
under.append(x)
else
over.append(x)
But it just seems like it should be possible to do this in a more elegant way knowing that the list is sorted. takewhile
and dropwhile
from itertools sound like the solution but then I would be iterating over the list twice.
Functionally, the best I can do is this:
i = 0
while sorted_list[i] < n:
i += 1
under = sorted_list[:i]
over = sorted_list[i:]
But I'm not even sure if it is actually better than just iterating over the list twice and it is definitely not more elegant.
I guess I'm looking for a way to get the list returned by takewhile
and the remaining list, perhaps, in a pair.
The correct solution here is the bisect
module. Use bisect.bisect
to find the index to the right of n
(or the index where it would be inserted if it's missing), then slice around that point:
import bisect # At top of file
split_idx = bisect.bisect(sorted_list, n)
under = sorted_list[:split_idx]
over = sorted_list[split_idx:]
While any solution is going to be O(n)
(you do have to copy the elements after all), the comparisons are typically more expensive than simple pointer copies (and associated reference count updates), and bisect
reduces the comparison work on a sorted list
to O(log n)
, so this will typically (on larger inputs) beat simply iterating and copying element by element until you find the split point.
Use bisect.bisect_left
(which finds the leftmost index of n
) instead of bisect.bisect
(equivalent to bisect.bisect_right
) if you want n
to end up in over
instead of under
.