Categorize list in Python
<p>What is the best way to categorize a list in python?</p>

<p>for example:</p>

<pre><code>totalist is below

totalist[1] = ['A','B','C','D','E']
totalist[2] = ['A','B','X','Y','Z']
totalist[3] = ['A','F','T','U','V']
totalist[4] = ['A','F','M','N','O']
</code></pre>

<p>Say I want to get the list where the first two items are <code>['A','B']</code>, basically <code>list[1]</code> and <code>list[2]</code>.  Is there an easy way to get these without iterate one item at a time?  Like something like this?</p>

<pre><code>if ['A','B'] in totalist
</code></pre>

<p>I know that doesn't work.</p>

<p>Just for fun, <code>itertools</code> solution to push per-element work to the C layer:</p>

<pre><code>from future_builtins import map  # Py2 only; not needed on Py3
from itertools import compress
from operator import itemgetter

# Generator
prefixes = map(itemgetter(slice(2)), totalist)
selectors = map(['A','B'].__eq__, prefixes)

# If you need them one at a time, just skip list wrapping and iterate
# compress output directly
matches = list(compress(totalist, selectors))
</code></pre>

<p>This could all be one-lined to:</p>

<pre><code>matches = list(compress(totalist, map(['A','B'].__eq__, map(itemgetter(slice(2)), totalist))))
</code></pre>

<p>but I wouldn't recommend it. Incidentally, if <code>totalist</code> might be a generator, not a re-iterable sequence, you'd want to use <code>itertools.tee</code> to double it, adding:</p>

<pre><code> totalist, forselection = itertools.tee(totalist, 2)
</code></pre>

<p>and changing the definition of <code>prefixes</code> to <code>map</code> over <code>forselection</code>, not <code>totalist</code>; since <code>compress</code> iterates both iterators in parallel, <code>tee</code> won't have meaningful memory overhead.</p>

<p>Of course, as others have noted, even moving to C, this is a linear algorithm. Ideally, you'd use something like a <code>collections.defaultdict(list)</code> to map from two element prefixes of each <code>list</code> (converted to <code>tuple</code> to make them legal <code>dict</code> keys) to a <code>list</code> of all <code>list</code>s with that prefix. Then, instead of linear search over N <code>list</code>s to find those with matching prefixes, you just do <code>totaldict['A', 'B']</code> and you get the results with <code>O(1)</code> lookup (and less fixed work too; no constant slicing).</p>

<p>Example precompute work:</p>

<pre><code>from collections import defaultdict

totaldict = defaultdict(list)
for x in totalist:
    totaldict[tuple(x[:2])].append(x)

# Optionally, to prevent autovivification later:
totaldict = dict(totaldict)
</code></pre>

<p>Then you can get <code>matches</code> effectively instantly for any two element prefix with just:</p>

<pre><code>matches = totaldict['A', 'B']
</code></pre>

python list