Retrieving lists of consecutive capitalised words from a list

Question

Ok, so given the string:

s = "Born in Honolulu Hawaii Obama is a graduate of Columbia University and Harvard Law School"

I want to retrieve:

[ ["Born"], ["Honolulu", "Hawaii", "Obama"], ["Columbia", "University"] ...]

Assuming that we have successfully tokenised the original string, my first psuedocodish attempt was:

def retrieve(tokens):
    results = []
    i = 0
    while i < len(tokens):
        if tokens[i][0].isupper():
            group = [tokens[i]]
            j = i + 1
            while i + j < len(tokens):
                if tokens[i + j][0].isupper():
                    group.append(tokens[i + j])
                    j += 1
                else:
                    break                
        i += 1
    return results

This is actually quite fast (well compared to some of my trying-to-be-pythonic attempts):

Timeit: 0.0160551071167 (1000 cycles)

Playing around with it, the quickest I can get is:

def retrive(tokens):
    results = []
    group = []
    for i in xrange(len(tokens)):
        if tokens[i][0].isupper():
            group.append(tokens[i])
        else:
            results.append(group)
            group = []
    results.append(group)
    return filter(None, results)

Timeit 0.0116229057312

Are there any more concise, pythonic ways to go about this (with similar execution times)?

Answer 1

A trivial optimisation that iterates on the tokens instead of by index (remember that in Python lists are iterables, it's unPythonic to iterate a list by index):

def retrieve(tokens):
    results = []
    group = []
    for token in tokens:
        if token[0].isupper():
            group.append(token)
        else:
            if group:  # group is not empty
                results.append(group)
            group = []  # reset group
    return results

A solution like @JeremyK's with a list comprehension and regular expressions is always going to be more compact. I am only giving this answer to point out how lists should be iterated.

Answer 2

Here is a solution using regular expressions

import re

def reMethod(pat, s):
    return [m.group().split() for m in re.finditer(pat, s)]

if __name__=='__main__':
    from timeit import Timer
    s = "Born in Honolulu Hawaii Obama is a graduate of Columbia University and Harvard Law School"
    pat = re.compile(r"([A-Z][a-z]*)(\s[A-Z][a-z]*)*")
    t1 = Timer(lambda: reMethod(pat, s))
    print "time:", t1.timeit(number=1000)
    print "returns:", reMethod(pat, s)

Ouptut:

time: 0.0183469604187
returns: [['Born'], ['Honolulu', 'Hawaii', 'Obama'], ['Columbia', 'University'], ['Harvard', 'Law', 'School']]

Answer 3

If you asked me, a pythonic solution would be the most readable solution. However in general, readable solutions are not always going to be the fastest. You can't have both, these are conflicting goals.

Pedro has covered what changes you can do to make it bit more pythonic so I won't repeat that. I can't think of anything else you can do to optimize that any further. However if you want a more purely pythonic solution (IMO), you can do this.

Use itertools.groupby. It groups consecutive items together with a common key into groups. Group them by whether they are capitalized and filter out the non-capitalized words.

This could be done as a one-liner but for readability, I'd write it like this:

from itertools import groupby

input_str = 'Born in Honolulu Hawaii Obama is a graduate of ' +
            'Columbia University and Harvard Law School'
words = input_str.split()
def is_capitalized(word):
    return bool(word) and word[0].isupper()
query = groupby(words, key=is_capitalized)
result = [list(g) for k, g in query if k]

This would run roughly twice as slow as your implementation, based on my tests on 1000 iterations.

0.0130150737235 # my implementation
0.0052368790507 # your implementation

You decide what your goals are.

Answer 4

I get a slight improvement by tweaking your retrive() function:

def retrieve1(tokens):
    results = []
    group = []
    for i in xrange(len(tokens)):
        if tokens[i][0].isupper():
            group.append(tokens[i])
        else:
            results.append(group)
            group = []
    results.append(group)
    return filter(None, results)

def retrieve2(tokens):
    results = []
    group = []
    for token in tokens:
        if token[0].isupper():
            group.append(token)
        elif group:
            results.append(group)
            group = []
    if group:
        results.append(group)
    return results

s = "Born in Honolulu Hawaii Obama is a graduate of Columbia University and Harvard Law School"

if __name__ == "__main__":
    import timeit
    print timeit.timeit('retrieve1(s.split(" "))', setup="from __main__ import (retrieve1, s)", number=100000)
    print timeit.timeit('retrieve2(s.split(" "))', setup="from __main__ import (retrieve2, s)", number=100000)

Results:

1.00138902664 (retrieve1)
0.744722127914 (retrieve2)

It's not necessary to iterate via xrange() as you can iterate tokens directly, and you need only append group to results if group is not an empty list -- no need to filter().