# Calculate co-occurrences of some words in documents dataset

It's a function to calculate the co-occurrences of some words in a news dataset, such as Techcrunch, Wired. The words list could be locations, products or people names.

Words list example:

["New York","Los Angeles","Chicago"]

Return result:

{"Chicago": {"New York" : 1, "Los Angeles": 2}}

The co-cocurrence between "Chicago" and "New York" is 1.

The problems:

The code below will caculate the same co-cocurrence of two words twice. And the time cost by a test dataset with 5 articles is 13.5s.

So for a dataset of 100k articles will cost about 75 hours. Is there any better solutions to improve the performance? Thanks!

"13.5 s ± 415 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)"

def get_co_occurrences(data):
com = defaultdict(dict)
for temp in data:
for i in range(len(city_list)-1):
for j in range(i+1, len(city_list)):
w1, w2 = city_list[i], city_list[j]
if " " + w1 + " " in temp and " " + w2 + " " in temp:
print(w1,w2)
if com[w1].get(w2) is None:
com[w1][w2] = 1
else:
com[w1][w2] += 1
return com


Edited: Python version

Python 3.6.2 | packaged by conda-forge | (default, Jul 23 2017, 22:59:30) [GCC 4.8.2 20140120 (Red Hat 4.8.2-15)] on linux

• I see 3 for loops embedded in each other, that's definitely your main problem. Please clarify your Python version as well. If this is Python 2, a different answer applies than Python 3. – Mast Jan 22 '18 at 16:51

1. There is no docstring. What does the function do? What argument does it take? What does it return?

2. city_list is a global variable. It would be better if this were passed as a parameter to the function: this would make it easier to reuse the function or write unit tests.

3. The algorithm has nothing to do with cities: it would work just as well for countries, or people, or any other set of search terms. It would make this generality clearer if city_list were named something like terms.

4. The variable names are vague and the code would be easier to understand if they were more specific. data is a collection of documents, so documents would be clearer. temp is a document from the collection, so document would be clearer. com is a data structure containing counts of co-occurrences that will be returned as the result of the function, so a name like result would be clearer.

5. There are two loops over i and j such that 0 <= i < j < len(city_list). These could be combined into a single loop using itertools.combinations:

for i, j in combinations(range(len(city_list)), 2):

6. The only purpose of these indexes is to select the two cities. It would simpler to iterate over the cities directly, avoiding the need for their indexes.

for w1, w2 in combinations(city_list, 2):

7. The data structure com is a mapping from first city to mappings from second cities to counts of co-occurrences. Unless you really need all those mappings it would simpler if the data structure were a mapping from pair of cities to count of co-occurrences.

8. When you have a data structure containing counts of items, using collections.Counter often makes the code simpler. First, create the counter object:

result = Counter()


and then increment the count:

if " " + w1 + " " in temp and " " + w2 + " " in temp:
result[w1, w2] += 1

9. Searching for terms with spaces before and after will miss matches at the beginnings and ends of documents. It would be more reliable to use a regular expression match together with the \b (word boundary) code:

if all(re.search(r'\b{}\b'.format(re.escape(term)), temp) for term in (w1, w2)):


(See the documentation for re.search and re.escape.)

10. The implementation in the post searches for every pair of cities in every document. If there are $k$ cities, $n$ documents, and $w$ words per document, then the overall runtime is $Θ(nk^2w)$.

An alternative approach is to check each word in each document to see if it is a city, and then to iterate over all pairs of cities found in the document. This has runtime $O(k + n(w \log w + \min(w, k)^2))$ since there can be $O(w)$ cities found in each document. This is more efficient than the original code in all cases, and much faster if (as one might expect) most words in the document are not search terms.

This could be implemented like this, by joining the search terms together into a single regular expression and then using the findall method:

from collections import Counter
from itertools import combinations
import re

def co_occurrences(documents, terms):
"""Return mapping from pairs of search terms to a count of documents
containing both terms.

"""
# Ensure that longer terms match in preference to shorter terms
# that happen to be initial substrings.
terms = sorted(terms, key=len, reverse=True)
# Regular expression matching any of the terms, as complete words.
search = re.compile(r'\b(?:{})\b'.format('|'.join(map(re.escape, terms))))
result = Counter()
for document in documents:
matches = sorted(set(search.findall(document)))
result.update(combinations(matches, 2))
return result

• Great great thanks for such long and detailed answer! – Tony Wang Jan 23 '18 at 7:30

Yes. You can do this in better way. You have to only find which city are in your data.

def get_co_occurrences(data):
com = defaultdict(dict)
for temp in data:
city_in_data = []
for city in city_list:
if city in temp:
city_in_data.append(city)
for cityA in city_in_data:
for cityB in city_in_data:
if cityB==cityA:
continue
if cityA not in com:
com[cityA] = {}
if cityB not in com[cityA]:
com[cityA][cityB] = 1
else:
com[cityA][cityB] += 1

return com