[python] sorted 함수와 key, 복잡한 정렬를 수행하는 법

CupRaccoon·2022년 7월 18일

* Python 3.8.2 버전 기준으로 작성되었습니다. 버전에 따라 제대로 동작하지 않을 수도 있습니다.

최근 모 기업 코딩테스트에 응시했는데 사전순을 포함한 여러 조건을 가진 정렬을 수행해야하는 문제를 만났다. 주어진 공식 Python Reference 외에는 참조가 불가능한 환경이어서 정렬을 구현하는데에 어려움을 겪었다.
그래서 코딩테스트가 끝난 후 python에서의 정렬과 sorted 함수에 대해서 자세히 알아봤고 한국어로 작성된 (가려운 곳을 긁어주는) 좋은 자료가 없는 것 같아 내가 나름대로 정리한 바를 작성하고자 한다.

1. sorted 함수

sorted(iterable, /, *, key=None, reverse=False)

sorted 함수는 iterable 객체를 정렬된 list 형태로 반환하는 함수이다.
argument로는 iterable, key, reverse를 가지고 있다.

iterable은 정렬에 사용될 객체를 argument로 받는데 iterable이라는 이름처럼 list, dictionary, set, string, tuple 등 iterable 객체를 모두 입력으로 받을 수 있다. (반환은 list로 이루어지는 점에 유의)

key는 iterable 객체의 각 요소에서 비교 키를 추출하는데 사용되는 argument가 하나인 함수를 입력으로 받는다. 기본값은 None으로 따로 지정하지 않았을 경우 iterable의 요소의 기본 비교연산을 이용해 정렬한다.

reverse는 Boolean 값으로 True일 경우 역순으로 정렬하고 기본값인 False일 경우 바로(일반적으로 오름차순) 정렬한다.

이제 글의 목적인 key를 기본적으로 활용하는 방식들과 key를 통해 어떻게 복잡한 정렬를 수행하는지 알아보자.

1) sorted 함수는 정렬에 대해서 stable을 보장한다. (정렬을 제외한 순서는 변경되지 않는다)
2) 생소한 /, *은 일종의 marker로 /는 / 앞의 argument가 positional-only argument임을 나타내고, *는 * 뒤의 argument가 keyword-only argument임을 나타낸다.
( What are positional-only arguments in Python?
: https://www.educative.io/answers/what-are-positional-only-arguments-in-python
python reference 용어집 중 argument 정의
: https://docs.python.org/ko/3.6/glossary.html )

2. key : key function

key는 argument가 하나이고 리턴 값도 하나인 함수를 지정해서 정렬을 수행할 수 있다. 정확히는 iterable 객체의 각 요소를 argument로 가지고 각각의 element에 대해서 "비교가능한" 값을 리턴해서 정렬에 사용한다.

예시를 위해서 몇 가지 iterable 객체를 생성하자.

foo1 = [[1,2,3,6,4,5,3],[1,2,7,2,3,3],[1,1,1,1,1],[1,4,1,1,1,1],[5,3,3,5,4]]

foo2 = [{'grade':'A','value':3},{'grade':'B','value':1},{'grade':'A','value':2},
        {'grade':'C','value':6},{'grade':'A','value':1}]
        
class Student:
    def __init__(self, grade, value):
        self.grade = grade
        self.value = value
    def __repr__(self):
         return repr((self.grade, self.value))
         
foo3 = [Student('A', 3),Student('B', 1),Student('A', 2),Student('C', 6),Student('A', 1)]

따라서 함수의 형태가 key에 적합하다면 기존의 함수를 이용할 수 있다.

sorted("This is a test string from Andrew".split(), key=str.lower)
>>>['a', 'Andrew', 'from', 'is', 'string', 'test', 'This']

sorted(foo1,key=sum)
>>>[[1, 1, 1, 1, 1], [1, 4, 1, 1, 1, 1], [1, 2, 7, 2, 3, 3], [5, 3, 3, 5, 4], [1, 2, 3, 6, 4, 5, 3]]

sorted(foo1,key=len)
>>>[[1, 1, 1, 1, 1], [5, 3, 3, 5, 4], [1, 2, 7, 2, 3, 3], [1, 4, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3]]

마찬가지로 사용자 정의 함수도 정렬에 사용할 수 있다.

def bar1(list):
    return list[0]+list[2]
    
sorted(foo1,key=bar1)
>>> [[1, 1, 1, 1, 1], [1, 4, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3], [1, 2, 7, 2, 3, 3], [5, 3, 3, 5, 4]]

다만 사용자 정의 함수를 따로 정의해서 사용하는 것보다 다음의 lambda식을 이용하는 것이 간편하다.

3. key : lambda

key를 사용하는 가장 일반적인 방법은 익명함수 lambda를 이용하는 것이다. lambda를 이용해서 객체의 특정 요소를 기준으로 정렬을 할 수 있다.

먼저 가장 간단하게 단일 index(,key, attribute,...)를 기준으로 정렬를 수행할 수 있다.

sorted(foo1,key=lambda x: x[1])
>>> [[1, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3], [1, 2, 7, 2, 3, 3], [5, 3, 3, 5, 4], [1, 4, 1, 1, 1, 1]]

sorted(foo2,key=lambda x: x['grade'])
>>> [{'grade': 'A', 'value': 3}, {'grade': 'A', 'value': 2}, {'grade': 'A', 'value': 1}, {'grade': 'B', 'value': 1}, {'grade': 'C', 'value': 6}]

sorted(foo3,key=lambda x: x.grade)
>>> [('A', 3), ('A', 2), ('A', 1), ('B', 1), ('C', 6)]

key를 지정하지 않았을 때와 마찬가지로 정수는 오름차순으로, 문자열은 사전순서대로 데이터 타입의 기존 비교 연산대로 정렬되는 것을 알 수 있다.
또한 grade가 'A'인 세 요소들이 정렬 전의 순서를 유지하는 것을 통해 정렬이 stable하다는 것도 확인할 수 있다.

2개 이상의 기준으로 정렬하고 싶은 경우 튜플 형태로 나타내면 된다. 튜플의 index가 낮은 순서대로 높은 우선순위를 가진다. ( multi-level sorting )

sorted(foo1,key=lambda x:(x[0],x[1]))
>>> [[1, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3], [1, 2, 7, 2, 3, 3], [1, 4, 1, 1, 1, 1], [5, 3, 3, 5, 4]]

sorted(foo1,key=lambda x:(len(x),x[1]))
>>> [[1, 1, 1, 1, 1], [5, 3, 3, 5, 4], [1, 2, 7, 2, 3, 3], [1, 4, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3]]

sorted(foo2,key=lambda x:(x['grade'],x['value']))
>>> [{'grade': 'A', 'value': 1}, {'grade': 'A', 'value': 2}, {'grade': 'A', 'value': 3}, {'grade': 'B', 'value': 1}, {'grade': 'C', 'value': 6}]

sorted(foo3,key=lambda x:(x.grade,x.value))
>>> [('A', 1), ('A', 2), ('A', 3), ('B', 1), ('C', 6)]

4. key : itemgetter(), attrgetter(), methodcaller()

3절의 방법이 매우 흔히 사용되기 때문에 python에서는 더 빠르고 쉽게 요소에 접근할 수 있도록 itemgetter(), attrgetter(), methodcaller()이라는 Operator Module Functions을 제공한다.

from operator import itemgetter, attrgetter

sorted(foo1,key=itemgetter(1))
>>> [[1, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3], [1, 2, 7, 2, 3, 3], [5, 3, 3, 5, 4], [1, 4, 1, 1, 1, 1]]

sorted(foo2,key=itemgetter('grade'))
>>> [{'grade': 'A', 'value': 3}, {'grade': 'A', 'value': 2}, {'grade': 'A', 'value': 1}, {'grade': 'B', 'value': 1}, {'grade': 'C', 'value': 6}]


sorted(foo3,key=attrgetter('grade'))
>>>[('A', 3), ('A', 2), ('A', 1), ('B', 1), ('C', 6)]

lambda식과 마찬가지로 다중 조건 정렬도 가능하다.

sorted(foo1,key=itemgetter(0,1))
>>> [[1, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3], [1, 2, 7, 2, 3, 3], [1, 4, 1, 1, 1, 1], [5, 3, 3, 5, 4]]


sorted(foo2,key=itemgetter('grade','value'))
>>> [{'grade': 'A', 'value': 1}, {'grade': 'A', 'value': 2}, {'grade': 'A', 'value': 3}, {'grade': 'B', 'value': 1}, {'grade': 'C', 'value': 6}]


sorted(foo3,key=attrgetter('grade','value'))
>>> [('A', 1), ('A', 2), ('A', 3), ('B', 1), ('C', 6)]

5. key : cmp_to_key

python에서는 C의 qsort()나 C++ sort()와 마찬가지로 compare 함수를 key로 사용할 수 있도록 cmp_to_key()라는 함수를 제공한다.

def compare(x,y):
    return x > y

위와 같은 형태로 x가 y보다 앞에 올 때 -1을 x, y간의 순서를 유지할 때는 0을 y가 x보다 앞에 올 때는 1을 리턴하도록 작성하면 된다. ( c/c++의 compare()와 순서가 반대임을 유의 )

def bar2(x, y):
    if x[0] < y[0]:
        return -1
    elif x[0] == y[0]:
        if x[1] < y[1]:
            return -1
        elif x[1] == y[1]:
            return 0
        else:
            return 1
    else:
        return 1
        
from functools import cmp_to_key

sorted(foo1, key=cmp_to_key(bar2))
>>> [[1, 1, 1, 1, 1], [1, 2, 3, 6, 4, 5, 3], [1, 2, 7, 2, 3, 3], [1, 4, 1, 1, 1, 1], [5, 3, 3, 5, 4]]

위와 같이 compare 함수를 이용해

sorted(foo1,key=lambda x:(x[0],x[1]))

위의 lambda식과 같은 동작을 하도록 구현할 수 있다. compare 함수를 직접 작성하는 특성상 조건이 많고 복잡한 정렬을 구현할 때 유리하다.

하지만 python reference Sorting HOW TO에서는 이 방법을 "The Old Way Using the cmp Parameter" 표현하고 있고 마찬가지로 cmp_to_key()의 설명에서도 "Transform an old-style comparison function to a key function" 이라고 표현하고 있으므로 python 2.x의 레거시 compare 함수를 사용해야 하는 것이 아니라면 python 3.x 에서는 사용을 지양하는 것이 바람직하다.

그렇다면 복잡한 정렬을 어떻게 구현할 수 있을까 ?

6. 복잡한 정렬

두 가지 종류의 복잡한 기준을 가정해서 앞서 정리한 방법들을 통해 어떻게 복잡한 정렬을 구현할 수 있는지 알아보자.
첫째로 (글을 쓴 계기가 된 코딩테스트의 문제를 간략화해서) 동적으로 입력받은 2차원 리스트를 인덱스 순서로 우선순위를 가지도록 하는 정렬
그리고 둘째로 3개의 key를 가진 dictionary를 요소로 가지는 list를 각 key에 대해서 내림차순, 오름차순, 내림차순 순서로 정렬을 구현해보자.

foo4 = [[2,1,3,6], [1, 2, 3 ,4], [1, 1, 1, 1], [2,1,3,3], [2,5,4,1]]
foo5 = [[4,3,2,5,1,2],[1,2,3,4,5,6],[4,3,2,5,1,1],[1,2,4,5,6,7],[2,1,1,1,1,1]]

foo6 = [{'grade': 'B', 'score': 81,'ID':201711}, {'grade': 'B', 'score': 95,'ID':201720}, {'grade': 'A', 'score': 99,'ID':201524},
        {'grade': 'A', 'score': 95, 'ID':201810}, {'grade': 'A', 'score': 95,'ID':201613}]

foo4, foo5를 동적으로 입력받는다고 가정하고 둘을 정렬해보자.

#lambda를 이용한 정렬
def bar3(input):
    ret = input
    for i in reversed(range(0,len(input[0]))):
        ret = sorted(ret,key=lambda x:x[i])
    return ret

bar3(foo4)
>>> [[1, 1, 1, 1], [1, 2, 3, 4], [2, 1, 3, 3], [2, 1, 3, 6], [2, 5, 4, 1]]
bar3(foo5)
>>> [[1, 2, 3, 4, 5, 6], [1, 2, 4, 5, 6, 7], [2, 1, 1, 1, 1, 1], [4, 3, 2, 5, 1, 1], [4, 3, 2, 5, 1, 2]]

#cmp_to_key()를 이용한 정렬
def bar4(list1,list2):
    for i in range(len(list1)):
        if list1[i] < list2[i]:
            return -1
        elif list1[i] == list2[i]:
            continue
        else:
            return 1
    return 1
    
sorted(foo4,key=cmp_to_key(bar4))
>>> [[1, 1, 1, 1], [1, 2, 3, 4], [2, 1, 3, 3], [2, 1, 3, 6], [2, 5, 4, 1]]
sorted(foo5,key=cmp_to_key(bar4))
>>> [[1, 2, 3, 4, 5, 6], [1, 2, 4, 5, 6, 7], [2, 1, 1, 1, 1, 1], [4, 3, 2, 5, 1, 1], [4, 3, 2, 5, 1, 2]]

이제 foo6를 score를 내림차순으로 grade를 오름차순으로 ID를 내림차순으로 하는 조건에 따라 정렬해보자.


#lambda를 이용한 정렬
def bar5(input):
    ret = input
    ret = sorted(ret,key=lambda x: x['ID'],reverse=True)
    ret = sorted(ret,key=lambda x: x['grade'],reverse=False)
    ret = sorted(ret,key=lambda x: x['score'],reverse=True)
    return ret

bar5(foo6)
>>> [{'grade': 'A', 'score': 99, 'ID': 201524}, {'grade': 'A', 'score': 95, 'ID': 201810}, {'grade': 'A', 'score': 95, 'ID': 201613}, {'grade': 'B', 'score': 95, 'ID': 201720}, {'grade': 'B', 'score': 81, 'ID': 201711}]

#cmp_to_key()를 이용한 정렬
def bar6(dic1,dic2):
    if dic1['score'] > dic2['score']:
        return -1
    elif dic1['score'] < dic2['score']:
        return 1
    else:
        if dic1['grade'] < dic2['grade']:
            return -1
        elif dic1['grade'] > dic2['grade']:
            return 1
        else:
            if dic1['ID'] > dic2['ID']:
                return -1
            elif dic1['ID'] < dic2['ID']:
                return 1
            else:
                return 0
                
sorted(foo6,key=cmp_to_key(bar6))
>>> [{'grade': 'A', 'score': 99, 'ID': 201524}, {'grade': 'A', 'score': 95, 'ID': 201810}, {'grade': 'A', 'score': 95, 'ID': 201613}, {'grade': 'B', 'score': 95, 'ID': 201720}, {'grade': 'B', 'score': 81, 'ID': 201711}]

코드를 살펴보면 lambda를 이용한 방법같은 경우 sorted 함수가 stable이 보장되는 것을 이용해서 우선순위의 역순으로 sorted 함수를 실행하면 의도한 대로 정렬할 수 있음을 알 수 있다. 이 때 sorted 함수의 수행횟수가 M번이라면 sorted 함수의 시간복잡도가 O(N·logN)이므로 전체 정렬의 시간복잡도는 O(M·N·logN)이 된다.

compare 함수를 이용한 방법같은 경우 전체 정렬의 시간복잡도는 compare 함수의 시간복잡도와 sorted 함수의 시간복잡도를 곱한 값임을 알 수 있다. 이 때 compare 함수의 시간복잡도는 O(M)이 됨을 직관적으로 알 수 있으므로 전체 정렬의 시간복잡도는 마찬가지로 O(M·N·logN)이 된다.

하지만 lambda를 사용한 경우 항상 M번 sorted 함수를 실행하지만 compare 함수같은 경우 worst case에서만 M번 연산을 수행하므로 정렬에 사용되는 데이터의 특성에 따라 더 빠르게 동작하는 것을 기대할 수 있다.

그런데 앞서 말했듯이 compare 함수는 사용을 지양해야 한다고 했다. 그렇다면 python에서 정렬에 비교 연산을 사용하는 방법이 없을까?

7. class and __lt__

compare 함수를 사용하는 대신 class를 이용해 데이터를 감싸고 비교 연산을 오버로딩해서 비슷하게 정렬을 구현할 수 있다. python에서는 class의 비교 연산을 위해 rich comparison이라고 불리는 비교 연산 메소드가 정의되어 있다. 각각 <, <=, ==, !=, >, >= 에 대응되는 __lt__, __le__, __eq__, __ne__, __gt__, __ge__의 6가지 메소드가 존재한다. PEP 8에서는 이 여섯 메소드를 모두 구현하는 것을 권고하지만 sorted 함수에서는 < 비교연산자만 사용하므로 __lt__만 구현해도 정렬은 가능하다.

x<y 는 x.__lt__(y)를 호출하므로 x<y 일 때 True를 x>=y일 때 False를 return하도록 작성하면 된다.

class ListClass:
    def __init__(self, list):
        self.list = list

    def __lt__(self, other):
        for i in range(len(self.list)):
            if self.list[i] < other.list[i]:
                return True
            elif self.list[i] == other.list[i]:
                continue
            else:
                return False
        return False

    def __repr__(self):
        return repr((self.list))

foo7 = [ListClass([4,3,2,5,1,2]),ListClass([1,2,3,4,5,6]),ListClass([4,3,2,5,1,1])
    ,ListClass([1,2,4,5,6,7]),ListClass([2,1,1,1,1,1])]

sorted(foo7)
>>> [[1, 2, 3, 4, 5, 6], [1, 2, 4, 5, 6, 7], [2, 1, 1, 1, 1, 1], [4, 3, 2, 5, 1, 1], [4, 3, 2, 5, 1, 2]]


class Student2:
    def __init__(self, grade, score, ID):
        self.grade = grade
        self.score = score
        self.ID = ID

    def __lt__(self,other):
        if self.score > other.score:
            return True
        elif self.score < other.score:
            return False
        else:
            if self.grade < other.grade:
                return True
            elif self.grade > other.grade:
                return False
            else:
                if self.ID > other.ID:
                    return True
                else:
                    return False
    def __repr__(self):
        return repr((self.grade, self.score, self.ID))



foo8 = [Student2('B', 81, 201711),Student2('B', 95, 201720),Student2('A', 99, 201524),
        Student2('A', 95, 201810),Student2('A' ,95, 201613)]
        
sorted(foo8)
>>> [('A', 99, 201524), ('A', 95, 201810), ('A', 95, 201613), ('B', 95, 201720), ('B', 81, 201711)]

정상적으로 정렬됨을 알 수 있다.

8. 결론

sorted 함수에서 key로 사용할 수 있는 lambda, operator module function, cmp_to_key와 class의 comparison method를 이용하는 방법을 알아봤다. 각각의 방법들이 특징이 다르기 때문에 상황에 맞게 사용하면 된다.

또한 cmp_to_key 함수를 사용하는 것을 지양해야 한다고 말하긴 했지만 코딩테스트와 같이 시간이 제한되어 있고 실제 프로덕션으로 이어지지 않는 코드라면 class를 작성하는 시간을 아낄 수 있으므로 경우에 따라 충분히 사용할 수 있다.

9. 참고자료

  1. python reference, Sorting HOW TO : https://docs.python.org/ko/3/howto/sorting.html
  2. PEP 207 – Rich Comparisons :
    https://peps.python.org/pep-0207/
  3. stackoverflow, Sort a list of lists with a custom compare function : https://stackoverflow.com/questions/5213033/sort-a-list-of-lists-with-a-custom-compare-function/57003713#57003713
  4. stackoverflow, How does Python's cmp_to_key function work? : https://stackoverflow.com/questions/16362744/how-does-pythons-cmp-to-key-function-work

여담

글을 작성하는 계기가 됬던 코딩테스트 문제는 "적절한" 조작을 거쳐서 문자열로 변환하면 쉽게 해결할 수 있었다고 한다......

profile
github : https://github.com/CupRaccoon

0개의 댓글