Coverage for o2/util/bit_mask_helper.py: 96%
99 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-16 11:18 +0000
1from functools import reduce
2from typing import Optional
5def bitmask_to_string(bitmask: int, pad_left: int = 24) -> str:
6 """Convert a bitmask to a string of 1s and 0s.
8 Pads the left side with 0s to the specified length.
9 """
10 return bin(bitmask)[2:].zfill(pad_left)
13def string_to_bitmask(bitmask: str) -> int:
14 """Convert a string of 1s and 0s to a bitmask."""
15 return int(bitmask, 2)
18def bitmask_to_array(bitmask: int, pad_left: int = 24) -> list[int]:
19 """Convert a bitmask to an array of integers."""
20 return [int(i) for i in bitmask_to_string(bitmask, pad_left)]
23def array_to_bitmask(bitmask: list[int]) -> int:
24 """Convert an array of integers to a bitmask."""
25 return string_to_bitmask("".join(map(str, bitmask)))
28def get_ranges_from_bitmask(bitmask: int) -> list[tuple[int, int]]:
29 """Get the ranges of 1s in a bitmask."""
30 bitmask_str = bitmask_to_string(bitmask)
31 ranges: list[tuple[int, int]] = []
32 start = None
33 for i in range(len(bitmask_str)):
34 if bitmask_str[i] == "1" and start is None:
35 start = i
36 elif bitmask_str[i] == "0" and start is not None:
37 ranges.append((start, i))
38 start = None
39 if start is not None:
40 ranges.append((start, len(bitmask_str)))
41 return ranges
44def has_overlap(bitmask_a: int, bitmask_b: int) -> bool:
45 """Check if two bitmasks have an overlap."""
46 return bitmask_a & bitmask_b != 0
49def any_has_overlap(bitmasks: list[int]) -> bool:
50 """Check if any of the bitmasks have an overlap with any other."""
51 acc = 0
52 for bitmask in bitmasks:
53 if acc & bitmask != 0:
54 return True
55 acc |= bitmask
56 return False
59def find_most_frequent_overlap(bitmasks: list[int], min_size: int = 1) -> Optional[tuple[int, int, int]]:
60 """Find the most frequent overlap in a list of bitmasks.
62 1. Convert bitmasks to bitarrays
63 2. Add all bitarrays together
64 3. Iterate over the resulting array and find the most frequent overlap of size >= min_size
65 4. Return start and end of the overlap
67 returns frequency of overlap, start, end; With [start, end) being the range of the overlap.
68 If no overlap is found, returns the longest range of 1s >= min_size (if any).
69 """
70 if len(bitmasks) == 0:
71 return None
72 bitarrays = [bitmask_to_array(bitmask) for bitmask in bitmasks]
73 summed_bitarray = reduce(lambda a, b: [x + y for x, y in zip(a, b)], bitarrays)
75 current_value = 0
76 current_start = 0
77 current_length = 0
79 max_start = None
80 max_end = None
81 max_value = 0
82 max_length = 0
84 current_sufficient_start = 0
85 current_sufficient_length = 0
87 max_sufficient_start = None
88 max_sufficient_end = None
89 max_sufficient_length = 0
91 for i, value in enumerate(summed_bitarray):
92 if value > 0:
93 if current_sufficient_length == 0:
94 current_sufficient_start = i
95 current_sufficient_length += 1
96 else:
97 if current_sufficient_length >= min_size and current_sufficient_length > max_sufficient_length:
98 max_sufficient_start = current_sufficient_start
99 max_sufficient_end = i
100 max_sufficient_length = current_sufficient_length
101 current_sufficient_start = i
102 current_sufficient_length = 0
104 if value == current_value:
105 current_length += 1
106 else:
107 # Check if the current range meets the requirements
108 if current_length >= min_size and (
109 (current_value == max_value and current_length > max_length) or (current_value > max_value)
110 ):
111 max_start = current_start
112 # End is exclusive
113 max_end = current_start + current_length
114 max_length = current_length
115 max_value = current_value
117 # Reset for new value
118 current_value = value
119 current_start = i
120 current_length = 1
122 # Final check after loop for sufficient length
123 if current_sufficient_length >= min_size and current_sufficient_length > max_sufficient_length:
124 max_sufficient_start = current_sufficient_start
125 max_sufficient_end = len(summed_bitarray)
126 max_sufficient_length = current_sufficient_length
128 # Final check after loop for max length
129 if current_length >= min_size and (
130 (current_value == max_value and current_length > max_length) or (current_value > max_value)
131 ):
132 max_start = current_start
133 # End is exclusive
134 max_end = current_start + current_length
136 if max_value > 0 and max_start is not None and max_end is not None:
137 return max_value, max_start, max_end
138 elif max_sufficient_start is not None and max_sufficient_end is not None:
139 return 1, max_sufficient_start, max_sufficient_end
140 else:
141 return None
144def find_mixed_ranges_in_bitmask(
145 bitmask: int,
146 length: int,
147 start_index: int,
148 max_start_index: int,
149 pad_left: int = 24,
150) -> list[tuple[int, int]]:
151 """In the given bitmask, find all (incl. overlapping) ranges of 1s and 0s.
153 The ranges need to be between start_index and max_start_index and contain
154 exactly `length` 1s. The 1s do NOT need to be consecutive. Meaning the sequence
155 will be at least length long, but due to intermediate 0s may be as long
156 as max_start_index-start_index.
158 NOTE: The returned range starts on the left, e.g., the left-most bit is 0.
159 The ranges are non-inclusive for the end index.
161 Args:
162 ----
163 bitmask (int): The input bitmask, treated as a binary sequence.
164 length (int): Exact number of `1`s required in the range.
165 start_index (int): The starting index of the search.
166 max_start_index (int): The maximum starting index of the search.
167 pad_left (int): The number of bits in the bitmask. Defaults to 24.
169 Returns:
170 -------
171 list[tuple[int, int]]: List of tuples where each tuple represents the start (inclusive) and end
172 (non-inclusive) indices of a valid range.
174 """
175 result = []
176 bitmask_length = pad_left
178 for i in range(start_index, min(max_start_index, bitmask_length - 1) + 1):
179 ones = 0
180 range_start = i
182 # Skip starting bits that are 0 -> we don't get 0-prefixed ranges
183 if not bitmask & (1 << (bitmask_length - i - 1)): # -1, because it's inclusive & 0-indexed
184 continue
186 for j in range(i, bitmask_length):
187 if bitmask & (1 << (bitmask_length - j - 1)): # Check if the bit is 1
188 ones += 1
190 if ones == length:
191 result.append((range_start, j + 1))
192 break
194 if ones > length:
195 break
197 return result