5
\$\begingroup\$

I have a set of strings that all consist of one or two letters followed by one or two digits. I need to sort by the letters then by the numeric value of the digits.

My solution is quite inefficient in terms of performance; how is it possible to improve it?

import re

unsorted_list = ['A1', 'A11', 'A12', 'A2', 'A3', 'B1', 'B12', 'EC1', 'EC21']
expected_result = ['A1', 'A2', 'A3', 'A11', 'A12', 'B1', 'B12', 'EC1', 'EC21']

unsorted_dict = {}
for item in unsorted_list:
    match = re.match(r"([A-Z]+)([0-9]+)", item, re.I)
    letters_in_item = match.groups()[0]
    numbers_in_item = int(match.groups()[1])
    if letters_in_item in unsorted_dict:
        unsorted_dict[letters_in_item].append(numbers_in_item)
    else:
        unsorted_dict[letters_in_item] = [numbers_in_item]

sorted_dict = dict(sorted(unsorted_dict.items()))

result = []
for key in sorted_dict:
    for value in sorted(sorted_dict[key]):
        result.append(key + str(value))

assert result == expected_result
\$\endgroup\$
0

2 Answers 2

8
\$\begingroup\$

Because we need to sort the data, the best we can do in terms of performance is to try and settle for \$O(N \log N\$).

What you have here is called natural sort. Apart from some libraries that you might use (e.g. natsort - which also handles more complex cases), I'm going to try to cleanup a bit your existing implementation.

NOTES:

  • Using a dictionary doesn't add any value or improvements to your code / algorithm but it rather adds overhead. As a rule of thumb, use dictionaries when you need to associate values with keys, so you can look them up efficiently (by key) later on;
  • You don't need to use sorted multiple times;
  • You could separate your code logic into different functions (benefits: easier to read, test, reuse, etc);
  • Type-hints are always welcome;

First, let's take out the regex you have:

COMPILED = re.compile(r"([A-Z]+)([0-9]+)", re.I)

For me, the biggest benefit to using re.compile() is being able to separate definition of the regex from its use. It might also offer some speed improvements but those cases are quire rare.

Now, what do we need to do?

  1. Split letters group from digits group (regex will take care of this);
  2. Sort by letters group then by digits group;
  3. Normalize data;

From what I can see you already know about sorted() but didn't use it to its fullest. Let's try and change that.

Having the first two points done (split + sort) is as simple as:

def nat_sort(items):
    # for me it feels more natural to use re.findall() instead of
    # re.match + re.group
    items_groups = [
        COMPILED.findall(item)[0]
        for item in items
    ]
    return sorted(items_groups, key=lambda group: (group[0], int(group[1])))

In the above, items_groups is going to look like this:

[('A', '1'), ('A', '11'), ('A', '12'), ('A', '2'), ('A', '3'), ('B', '1'), ('B', '12'), ('EC', '1'), ('EC', '21')]

The next line is just going to sort our items by the first group (letters), then by the last one (digits), giving us the following:

[('A', '1'), ('A', '2'), ('A', '3'), ('A', '11'), ('A', '12'), ('B', '1'), ('B', '12'), ('EC', '1'), ('EC', '21')]

items_groups could also be sorted in-place to avoid creating an extra list:

def nat_sort(items):
    items_groups = [
        COMPILED.findall(item)[0]
        for item in items
    ]
    items_groups.sort(key=lambda group: (group[0], int(group[1])))
    return items_groups

The remaining thing we have to do is to transform the data we have into our desired format:

def main():
    items = ['A1', 'A11', 'A12', 'A2', 'A3', 'B1', 'B12', 'EC1', 'EC21']
    return [''.join(item) for item in nat_sort(items)]

Full code:

import re

COMPILED = re.compile(r"([A-Z]+)([0-9]+)", re.I)


def nat_sort(items):
    items_groups = [
        COMPILED.findall(item)[0]
        for item in items
    ]
    items_groups.sort(key=lambda group: (group[0], int(group[1])))
    return items_groups


def main():
    items = ['A1', 'A11', 'A12', 'A2', 'A3', 'B1', 'B12', 'EC1', 'EC21']
    return [''.join(item) for item in nat_sort(items)]


if __name__ == '__main__':
    print(main())

Or as recommended in the comments, have everything contained within the same function (BONUS: type-hints included):

import re
from typing import List

COMPILED = re.compile(r"([A-Z]+)([0-9]+)", re.I)


def nat_sort(items: List[str]) -> List[str]:
    items_groups = [
        COMPILED.findall(item)[0]
        for item in items
    ]
    items_groups.sort(key=lambda group: (group[0], int(group[1])))
    return [''.join(item) for item in items_groups]


def main() -> List[str]:
    items = ['A1', 'A11', 'A12', 'A2', 'A3', 'B1', 'B12', 'EC1', 'EC21']
    return nat_sort(items)


if __name__ == '__main__':
    print(main())

Not sure if this comes with an improvement in terms of complexity (my complexity analysis is quite rusty) but my suggestion has \$O(N \log N\$) time complexity and \$O(N)\$ space complexity.

There are probably better, easier-to-read, more intuitive ways of doing this but this felt natural to me.

\$\endgroup\$
2
  • 2
    \$\begingroup\$ findall()[0] seems rather strange to me ... Also nat_sort() should return the input type, thus the ''.join() should be included in nat_sort(). \$\endgroup\$
    – stefan
    Mar 11, 2022 at 12:36
  • \$\begingroup\$ Thanks for your kind help, explanation and the time you put into improving mine with explanations, Alex! The one here is glorious: stackoverflow.com/a/16090640/509977 \$\endgroup\$
    – Pitto
    Mar 11, 2022 at 21:28
2
\$\begingroup\$

match.groups()[0] returns all the groups of the match, and then indexes into the returned tuple to extract the first one. match.groups()[1] repeats this to return the second. There is no need to return the entire match.groups() tuple; you can simply and more efficiently use match.group(1) and match.group(2). For more compact code, use match[1] and match[2].

Defining a function that returns match[1], int(match[2]) gives the required sort key. You can just use this as the key= argument to sorted() for a very short & compact solution:

import re

COMPILED = re.compile(r"([A-Z]+)([0-9]+)", re.I)

def sort_key(item: str) -> tuple[str, int]:
    match = COMPILED.match(item)
    return match[1], int(match[2])

unsorted_list = ['A1', 'A11', 'A12', 'A2', 'A3', 'B1', 'B12', 'EC1', 'EC21']
expected_result = ['A1', 'A2', 'A3', 'A11', 'A12', 'B1', 'B12', 'EC1', 'EC21']

result = sorted(unsorted_list, key=sort_key)

assert result == expected_result
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.