如何在O(k*N +k*Q)复杂度下高效统计K-mer并按查询顺序输出？

Looks like your code is hitting timeouts because of a couple of inefficient operations—let's break down the issues and fix them to meet the required time complexity of O(kN + kQ).

What's Causing the Timeout? #

Your core logic is on the right track, but two key bottlenecks are slowing things down:

1. Checking membership in a list: When you do if current in mykmer, mykmer is a list, and this operation takes O(Q) time per check. With up to O(N) k-mers to verify, this pushes your total time complexity to O(kN + NQ), which gets way too slow when Q is large.
2. Inefficient DNA string concatenation: Building the dna string with repeated dna += ... creates new string objects every time (since Python strings are immutable). For large DNA sequences, this adds unnecessary overhead.

Optimized Solution #

Here's the revised code that fixes these issues while keeping your original logic intact:

import sys

def main():
    # Read all input at once for efficiency
    all_lines = sys.stdin.read().splitlines()
    line_ptr = 0
    
    # Parse parameters
    l, n, k, q = map(int, all_lines[line_ptr].split())
    line_ptr += 1
    
    # Read DNA sequence efficiently
    dna_parts = []
    for _ in range(l):
        dna_parts.append(all_lines[line_ptr].strip())
        line_ptr += 1
    dna = ''.join(dna_parts)
    
    # Prepare query storage: list for order, dict for counts, set for O(1) lookups
    query_list = []
    count_dict = {}
    query_set = set()
    
    for _ in range(q):
        kmer = all_lines[line_ptr].strip()
        line_ptr += 1
        query_list.append(kmer)
        count_dict[kmer] = 0
        query_set.add(kmer)
    
    # Handle edge case: k is longer than DNA length
    if k > len(dna):
        for kmer in query_list:
            print(f"{kmer} {count_dict[kmer]}")
        return
    
    # Sliding window to count relevant k-mers
    current_kmer = dna[:k]
    if current_kmer in query_set:
        count_dict[current_kmer] += 1
    
    for idx in range(k, len(dna)):
        # Update k-mer by sliding the window (O(k) per update, total O(k*N))
        current_kmer = current_kmer[1:] + dna[idx]
        if current_kmer in query_set:
            count_dict[current_kmer] += 1
    
    # Output results in query order
    for kmer in query_list:
        print(f"{kmer} {count_dict[kmer]}")

if __name__ == "__main__":
    main()

Key Optimizations Explained #

1. O(1) Membership Checks with a Set:

   • We store all query k-mers in a set (query_set). Checking if a generated k-mer is in the set takes O(1) time instead of O(Q), which brings the DNA processing phase down to O(k*N) time.

2. Efficient DNA String Construction:

   • Instead of concatenating strings directly, we collect DNA segments in a list and use ''.join() at the end. This avoids creating intermediate string objects and reduces the time to build the DNA string from O(l*L) (where L is average line length) to O(N).

3. Cleaner Sliding Window Logic:

   • We initialize the first k-mer upfront, then iterate from index k onwards to slide the window. This removes the unnecessary if j != k-1 check, making the code simpler and slightly faster.

4. Edge Case Handling:

   • We add a check for when k is longer than the DNA sequence length, which immediately outputs 0 for all queries without unnecessary processing.

Time Complexity Verification #

• Query Processing: Each query k-mer is added to the list, dict, and set in O(k) time (since we process each character of the k-mer), so total O(k*Q).
• DNA Processing: Generating each k-mer takes O(k) time, and there are O(N) k-mers, so total O(k*N).
• Output: Printing each result takes O(k) time per query, total O(k*Q).

This adds up to the required O(kN + kQ) time complexity, which will handle large inputs without timing out.

内容的提问来源于stack exchange，提问作者Waiwai Kim