숫자의 신

백준 1422번: 숫자의 신

아이디어

a와 b 두 숫자중에 어떤 숫자가 앞에 와야 하는지 생각해보자. 우선 무조건 앞자리가 큰 숫자가 맨 앞에 와야 한다. 문제는 앞부분이 똑같은 두개의 숫자(232와 23, 212와 21)를 비교하는 경우이다. 긴 숫자를 a, 짧은 숫자를 b라고 할 때 a에서 겹치는 부분을 제거한 substirng과 b를 비교한다. 이를 계속해서 반복하면 어떤 숫자를 앞에 두면 좋을지 결정할 수 있다. 232와 23중 어떤 숫자가 앞에 와야 할지 직접 계산해보자.
232 vs 23
2 vs 23
2 vs 3
23이 앞에 와야 한다.
또 N이 K보다 큰 경우에는 그냥 가장 큰 수를 여러번 사용하면 된다. 뭔가 앞자리가 큰 숫자를 써야 할 것 같지만 착각이다. 당연히 같은 숫자는 연속해서 사용해야한다.
ex) 9, 80 -> 98080(O) 80980(X)

코드

import sys
from functools import cmp_to_key
input = sys.stdin.readline

K, N = map(int, input().split())
numlist = [input().rstrip() for _ in range(K)]

def compare(a, b):
    if a == b:
        return 0
    if len(a) >= len(b):
        for i in range(len(b)):
            if a[i] > b[i]:
                return -1
            elif a[i] < b[i]:
                return 1
        return compare(a[len(b):], b)
    else:
        for i in range(len(a)):
            if a[i] > b[i]:
                return -1
            elif a[i] < b[i]:
                return 1
        return compare(a, b[len(a):])

s_list = sorted(numlist, key=cmp_to_key(compare))
m = ''
for i in s_list:
    if len(i) > len(m):
        m = i
    elif int(i) > int(m):
        m = i

cnt = N - K
for i in s_list:
    if i == m:
        while cnt > 0:
            print(i, end='')
            cnt -= 1
    print(i, end='')

설명

functools라는 생소한 라이브러리를 사용했다.
기본적으로 Python3에서 정렬 기준을 정하는 key에는 argument 하나를 받는 함수를 사용한다.

>>> student_tuples = [
...     ('john', 'A', 15),
...     ('jane', 'B', 12),
...     ('dave', 'B', 10),
... ]
>>> sorted(student_tuples, key=lambda student: student[2])   # sort by age
[('dave', 'B', 10), ('jane', 'B', 12), ('john', 'A', 15)]

이런 식으로 key에 써넣은 함수가 반환하는 값을 기준으로 정렬한다. 하지만 functools의 cmp_to_key를 사용하면 구식 비교 함수를 key 함수로 변환할 수 있다. 구식 비교 함수가 무엇일까?
원소 두 개(a, b)를 인자로 받아서 함수가 음수를 반환하면 a가 b보다 앞에 오게 정렬하고 양수를 반환하면 b가 a보다 앞에 오게 정렬하는 함수다. 약간 자바에서 compareTo 오버라이딩 하는 그런 느낌?
Python2에서는 이런 식으로 사용했다

sorted(mylist, cmp=lambda item1, item2: fitness(item1) - fitness(item2))

Python3에서는 이런 식으로 사용한다.

from functools import cmp_to_key
sorted(mylist, key=cmp_to_key(lambda item1, item2: fitness(item1) - fitness(item2)))

비교 함수가 반환하는 값을 기준으로 sort해준다.

여담

생각해봤는데 비교 함수 작성할 때 그냥 a랑 b 숫자 두개 이어붙여서 크기를 비교하면 훨씬 간단하게 풀 수 있지 않았을까? 나는 붕어다.

참고

Sort a list of lists with a custom compare function

functools