Python search one million strings in a file and count occurrences of each string

Question

This is more about to find the fastest way to do it. I have a file1 which contains about one million strings(length 6-40) in separate line. I want to search each of them in another file2 which contains about 80,000 strings and count occurrence(if small string is found in one string multiple times, the occurence of this string is still 1). If anyone is interested to compare performance, there is link to download file1 and file2. dropbox.com/sh/oj62918p83h8kus/sY2WejWmhu?m

What i am doing now is construct a dictionary for file 2, use strings ID as key and string as value. (because strings in file2 have duplicate values, only string ID is unique) my code is

for line in file1:
   substring=line[:-1].split("\t")
   for ID in dictionary.keys():
       bigstring=dictionary[ID]
       IDlist=[]
       if bigstring.find(substring)!=-1:
           IDlist.append(ID)
   output.write("%s\t%s\n" % (substring,str(len(IDlist))))

My code will take hours to finish. Can anyone suggest a faster way to do it? both file1 and file2 are just around 50M, my pc have 8G memory, you can use as much memory as you need to make it faster. Any method that can finish in one hour is acceptable:)

Here, after I have tried some suggestions from these comments below, see performance comparison, first comes the code then it is the run time.

Some improvements suggested by Mark Amery and other peoples

import sys
from Bio import SeqIO

#first I load strings in file2 to a dictionary called var_seq, 
var_seq={}
handle=SeqIO.parse(file2,'fasta')
for record in handle:
    var_seq[record.id]=str(record.seq)

print len(var_seq) #Here print out 76827, which is the right number. loading file2 to var_seq doesn't take long, about 1 second, you shall not focus here to improve performance

output=open(outputfilename,'w')
icount=0
input1=open(file1,'r')
for line in input1:
    icount+=1
    row=line[:-1].split("\t")
    ensp=row[0]   #ensp is just peptides iD
    peptide=row[1] #peptides is the substrings i want to search in file2
    num=0
    for ID,bigstring in var_seq.iteritems(): 
        if peptide in bigstring:
            num+=1

    newline="%s\t%s\t%s\n" % (ensp,peptide,str(num))
    output.write(newline)
    if icount%1000==0:
        break

input1.close()
handle.close()
output.close()

It will take 1m4s to finish. Improved 20s compared to my old one

#######NEXT METHOD suggested by entropy

from collections import defaultdict
var_seq=defaultdict(int)
handle=SeqIO.parse(file2,'fasta')
for record in handle:
    var_seq[str(record.seq)]+=1

print len(var_seq) # here print out 59502, duplicates are removed, but occurances of duplicates are stored as value 
handle.close()

output=open(outputfilename,'w')
icount=0

with open(file1) as fd:
    for line in fd:
        icount+=1
        row=line[:-1].split("\t")
        ensp=row[0]
        peptide=row[1]
        num=0
        for varseq,num_occurrences in var_seq.items():
            if peptide in varseq:
                num+=num_occurrences

    newline="%s\t%s\t%s\n" % (ensp,peptide,str(num))
    output.write(newline)
    if icount%1000==0:
        break

output.close()

This one takes 1m10s,not faster as expected since it avoids searching duplicates, don't understand why.

Haystack and Needle method suggested by Mark Amery, which turned out to be the fastest, The problem of this method is that counting result for all substrings is 0, which I don't understand yet.

Here is the code I implemented his method.

class Node(object):
    def __init__(self):
        self.words = set()
        self.links = {}

base = Node()

def search_haystack_tree(needle):
    current_node = base
    for char in needle:
        try:
            current_node = current_node.links[char]
        except KeyError:
            return 0
    return len(current_node.words)

input1=open(file1,'r')
needles={}
for line in input1:
    row=line[:-1].split("\t")
    needles[row[1]]=row[0]

print len(needles)

handle=SeqIO.parse(file2,'fasta')
haystacks={}
for record in handle:
    haystacks[record.id]=str(record.seq)

print len(haystacks)

for haystack_id, haystack in haystacks.iteritems(): #should be the same as enumerate(list)
    for i in xrange(len(haystack)):
        current_node = base
        for char in haystack[i:]:
            current_node = current_node.links.setdefault(char, Node())
            current_node.words.add(haystack_id)

icount=0
output=open(outputfilename,'w')
for needle in needles:
    icount+=1
    count = search_haystack_tree(needle)
    newline="%s\t%s\t%s\n" % (needles[needle],needle,str(count))
    output.write(newline)
    if icount%1000==0:
        break

input1.close()
handle.close()
output.close()

It takes only 0m11s to finish, which is much faster than other methods. However, I don't know it is my mistakes to make all counting result as 0, or there is a flaw in the Mark's method.

Answer 1

Your code doesn't seem like it works(are you sure you didn't just quote it from memory instead of pasting the actual code?)

For example, this line:

substring=line[:-1].split("\t")

will cause substring t be a list. But later you do:

if bigstring.find(substring)!=-1:

That would cause an error if you call str.find(list).

In any case, you are building lists uselessly in your innermost loop. This:

IDlist=[]
if bigstring.find(substring)!=-1:
     IDlist.append(ID)

 #later
 len(IDlist)

That will uselessly allocate and free lists which would cause memory thrashing as well as uselessly bogging everything down.

This is code that should work and uses more efficient means to do the counting:

from collections import defaultdict

dictionary = defaultdict(int)
with open(file2) as fd:
    for line in fd:
        for s in line.split("\t"):
            dictionary[s.strip()] += 1

with open(file1) as fd:
    for line in fd:
        for substring in line.split('\t'):
            count = 0
            for bigstring,num_occurrences in dictionary.items():
                if substring in bigstring:
                    count += num_occurrences
            print substring, count

PS: I am assuming that you have multiple words per line that are tab-split because you do line.split("\t") at some point. If that is wrong, it should be easy to revise the code.

PPS: If this ends up being too slow for your use(you'd have to try it, but my guess is this should run in ~10min given the number of strings you said you had). You'll have to use suffix trees as one of the comments suggested.

Edit: Amended the code so that it handles multiple occurrences of the same string in file2 without negatively affecting performance

Edit 2: Trading maximum space for time.

Below is code that will consume quite a bit of memory and take a while to build the dictionary. However, once that's done, each search out of the million strings to search for should complete in the time it takes for a single hashtable lookup, that is O(1).

Note, I have added some statements to log the time it takes for each step of the process. You should keep those so you know which part of the time is taken when searching. Since you are testing with 1000 strings only this matters a lot since if 90% of the cost is the build time, not the search time, then when you test with 1M strings you will still only be doing that once, so it won't matter

Also note that I have amended my code to parse file1 and file2 as you do, so you should be able to just plug this in and test it:

from Bio import SeqIO
from collections import defaultdict
from datetime import datetime

def all_substrings(s):
    result = set()
    for length in range(1,len(s)+1):
        for offset in range(len(s)-length+1):
            result.add(s[offset:offset+length])
    return result

print "Building dictionary...."
build_start = datetime.now()

dictionary = defaultdict(int)
handle = SeqIO.parse(file2, 'fasta')
for record in handle:
    for sub in all_substrings(str(record.seq).strip()):
        dictionary[sub] += 1
build_end = datetime.now()
print "Dictionary built in: %gs" % (build-end-build_start).total_seconds()

print "Searching...\n"
search_start = datetime.now()

with open(file1) as fd:
    for line in fd:
        substring = line.strip().split("\t")[1]
        count = dictionary[substring]
        print substring, count

search_end = datetime.now()
print "Search done in: %gs" % (search_end-search_start).total_seconds()

Answer 2

I'm not an algorithms whiz, but I reckon this should give you a healthy performance boost. You need to set 'haystacks' to be a list of the big words you want to look in, and 'needles' to be a list of the substrings you're looking for (either can contain duplicates), which I'll let you implement on your end. It'd be great if you could post your list of needles and list of haystacks so that we can easily compare performance of proposed solutions.

haystacks = <some list of words>
needles = <some list of words>

class Node(object):
    def __init__(self):
        self.words = set()
        self.links = {}

base = Node()

for haystack_id, haystack in enumerate(haystacks):
    for i in xrange(len(haystack)):
        current_node = base
        for char in haystack[i:]:
            current_node = current_node.links.setdefault(char, Node())
            current_node.words.add(haystack_id)

def search_haystack_tree(needle):
    current_node = base
    for char in needle:
        try:
            current_node = current_node.links[char]
        except KeyError:
            return 0
    return len(current_node.words)

for needle in needles:
    count = search_haystack_tree(needle)
    print "Count for %s: %i" % (needle, count)

You can probably figure out what's going on by looking at the code, but just to put it in words: I construct a huge tree of substrings of the haystack words, such that given any needle, you can navigate the tree character by character and end up at a node which has attached to it the set of all haystack ids of haystacks containing that substring. Then for each needle we just go through the tree and count the size of the set at the end.