Оптимизация решения, преобразование данных из строки в словарь и обратно

есть задание

коротко суть:

дана строка вида

"""John Daggett, 341 King Road, Plymouth MA
Alice Ford, 22 East Broadway, Richmond VA
Sal Carpenter, 73 6th Street, Boston MA"""

Необходимо преобразовать в строку вида:

"Massachusetts\r\n..... John Daggett 341 King Road Plymouth Massachusetts\r\n..... Sal Carpenter 73 6th Street Boston Massachusetts\r\n Virginia\r\n..... Alice Ford 22 East Broadway Richmond Virginia"

Где все значения сгруппированы по штатам, штаты начинаются с новой строки и отделены символом перевода каретки.

При этом, штаты должны быть в алфавитном порядке и значения в штате тоже должны быть в алфавитном порядке.

Мое решение:

def by_state(str):
    d = {'AZ': 'Arizona',
     'CA': 'California',
     'ID': 'Idaho',
     'IN': 'Indiana',
     'MA': 'Massachusetts',
     'OK': 'Oklahoma',
     'PA': 'Pennsylvania',
     'VA': 'Virginia'}
    
    # print(str)
    
    sep = sorted([[j for j in i.replace(',', '').split(' ')][:-1] + [d[k] for k in i.split() if k in d] for i in str.split('\n')]) 
   # здесь я разбиваю большую строку на список строк и сразу меняю аббревиатуру штата на название и сортирую это все по первому символу первого значения
    
    dct = {}
    for i in sep:
        if i[-1] in dct:
            dct[i[-1]] += '\r\n..... ' + ' '.join(i)
        else:
            dct[i[-1]] = '\r\n..... ' + ' '.join(i)
   # затем создаю словарь, в котором ключи - названия штатов, значения - строки
    
    ans = []
    for i in sorted(dct):
        ans += [f' {i}{dct[i]}']
        
    return '\r\n'.join(ans)[1:]
    # здесь я извращаюсь с выводом, что бы получить из словаря строку, нужного вида.

первое: решение громоздкое и ужасное, но как его оптимизировать

второе: недружественный проверяльщик сервиса выдает ошибку вида:

STDERR
Traceback (most recent call last):
  File "/workspace/default/tests.py", line 94, in <module>
    test.assert_equals(by_state(s), sol)
  File "/workspace/default/solution.py", line 17, in by_state
    if i[-1] in dct:
IndexError: list index out of range

хотя у меня в pyCharm все ок...

Кака оптимизировать это решение?


Ответы (2 шт):

Автор решения: kristal

Мой вариант решения:

def by_state(raw_string):
    state_code = {'AZ': 'Arizona',
                  'CA': 'California',
                  'ID': 'Idaho',
                  'IN': 'Indiana',
                  'MA': 'Massachusetts',
                  'OK': 'Oklahoma',
                  'PA': 'Pennsylvania',
                  'VA': 'Virginia'}
    result_string = last_state = ""
    split_string = map(lambda x: x.strip().split(", "), raw_string.split("\n"))
    processed_list = [item[:-1] + item[-1].split() for item in split_string if item[0]]
    for line in sorted(processed_list, key=lambda address: (address[-1], address[0], address[1])):
        line[-1] = full_state_name = state_code[line[-1]]
        if last_state != full_state_name:
            last_state = full_state_name
            result_string += f" {last_state}\r\n"
        result_string += f"..... {' '.join(line)}\r\n"
    return result_string.strip()

Можно чутка сократить если использовать регулярки:

import re

def by_state(raw_string):
    state_code = {'AZ': 'Arizona',
                  'CA': 'California',
                  'ID': 'Idaho',
                  'IN': 'Indiana',
                  'MA': 'Massachusetts',
                  'OK': 'Oklahoma',
                  'PA': 'Pennsylvania',
                  'VA': 'Virginia'}
    result_string = last_state = ""
    processed_list = re.findall(r"(.+?), (.+?), (.+?) ([A-Z]{2})", raw_string)
    for name, addr, city, state in sorted(processed_list, key=lambda address: (address[-1], address[0], address[1])):
        full_state_name = state_code[state]
        if last_state != full_state_name:
            last_state = full_state_name
            result_string += f" {last_state}\r\n"
        result_string += f"..... {name} {addr} {city} {full_state_name}\r\n"
    return result_string.strip()
→ Ссылка
Автор решения: Alex Titov
from itertools import groupby

d = {'AZ': 'Arizona',
     'CA': 'California',
     'ID': 'Idaho',
     'IN': 'Indiana',
     'MA': 'Massachusetts',
     'OK': 'Oklahoma',
     'PA': 'Pennsylvania',
     'VA': 'Virginia'}

def by_state(str):
    # Почти то же самое, но сразу сортировка по штатам и по имени
    sep = sorted([[j for j in i.replace(',', '').split(' ')][:-1] + [d[k] for k in i.split() if k in d] for i in
                  str.split('\n')], key=lambda x: x[-1:] + x[:-1] )
    # Отдельной строкой группировка, но можно и в одну строку
    gb = groupby(sep, key=lambda x: x[-1])
    return '\r\n '.join(key +"\r\n..... " + "\r\n..... ".join(" ".join(s) for s in group) for key, group in gb)

P.S. А рандомные тесты не проходит, что-то в разборе строки не так, возможно.

P.P.S Похоже, в тестах встречаются пустые строки, чуть упростил разбор и поставил фильтр от пустых строк, чтобы проходило

def by_state(str):
    lst = [s.rsplit(maxsplit=1) for s in str.replace(',', '').split('\n')]
    sep = sorted(([s[0], d[s[1]]] for s in lst if s), key=lambda x: (x[1], x[0]))
    gb = groupby(sep, key=lambda x: x[1])
    return '\r\n '.join(key +"\r\n..... " + "\r\n..... ".join(" ".join(s) for s in group) for key, group in gb)
→ Ссылка