threads 2

Author: kiri11

rounds/3_dna/solution.py

@@ -16,16 +16,17 @@ def find_matches(fasta_path: str, pattern: bytes) -> list[tuple[str, list[int]]]
     if not pattern:
         return []
 
-    pattern_values = np.frombuffer(pattern, dtype=np.uint8)
     pattern_len = len(pattern)
+    pattern_prefix = np.frombuffer(pattern[:4], dtype=np.uint32)[0]
+    pattern_suffix = np.frombuffer(pattern[4:], dtype=np.uint32)[0]
 
     with open(fasta_path, "rb") as file:
         data = file.read()
 
     records = data.split(b">")[1:]
     worker_count = min(_MAX_WORKERS, os.cpu_count() or 1, len(records))
     if worker_count <= 1:
-        return _scan_records(records, pattern_values, pattern_len)
+        return _scan_records(records, pattern_prefix, pattern_suffix, pattern_len)
 
     chunk_size = (len(records) + worker_count - 1) // worker_count
     chunks = [
@@ -36,43 +37,58 @@ def find_matches(fasta_path: str, pattern: bytes) -> list[tuple[str, list[int]]]
         groups = executor.map(
             _scan_records,
             chunks,
-            [pattern_values] * len(chunks),
+            [pattern_prefix] * len(chunks),
+            [pattern_suffix] * len(chunks),
             [pattern_len] * len(chunks),
         )
 
     return [match for group in groups for match in group]
 
 
 def _scan_records(
-    records: list[bytes], pattern_values: np.ndarray, pattern_len: int
+    records: list[bytes],
+    pattern_prefix: np.uint32,
+    pattern_suffix: np.uint32,
+    pattern_len: int,
 ) -> list[tuple[str, list[int]]]:
     matches: list[tuple[str, list[int]]] = []
     for record in records:
-        match = _scan_record(record, pattern_values, pattern_len)
+        match = _scan_record(record, pattern_prefix, pattern_suffix, pattern_len)
         if match is not None:
             matches.append(match)
     return matches
 
 
 def _scan_record(
-    record: bytes, pattern_values: np.ndarray, pattern_len: int
+    record: bytes,
+    pattern_prefix: np.uint32,
+    pattern_suffix: np.uint32,
+    pattern_len: int,
 ) -> tuple[str, list[int]] | None:
     record_id, _, wrapped_sequence = record.partition(_NEWLINE)
     sequence = wrapped_sequence.replace(_NEWLINE, b"")
     sequence_len = len(sequence)
     if sequence_len < pattern_len:
         return None
 
-    sequence_values = np.frombuffer(sequence, dtype=np.uint8)
     candidate_count = sequence_len - pattern_len + 1
-    positions_mask = sequence_values[:candidate_count] == pattern_values[0]
-    for pattern_index in range(1, pattern_len):
-        positions_mask &= (
-            sequence_values[pattern_index : candidate_count + pattern_index]
-            == pattern_values[pattern_index]
-        )
+    prefixes = np.ndarray(
+        shape=(candidate_count,),
+        dtype=np.uint32,
+        buffer=sequence,
+        strides=(1,),
+    )
+    candidates = np.nonzero(prefixes == pattern_prefix)[0]
+    if not candidates.size:
+        return None
 
-    positions = np.nonzero(positions_mask)[0]
+    suffixes = np.ndarray(
+        shape=(candidate_count,),
+        dtype=np.uint32,
+        buffer=memoryview(sequence)[4:],
+        strides=(1,),
+    )
+    positions = candidates[suffixes[candidates] == pattern_suffix]
     if positions.size:
         return record_id.decode("ascii"), positions.tolist()
     return None