Python을 사용하면서, 코딩테스트에 도움이 되는 유용한 함수/기능들을 정리하려고 한다.
프로그래머스의 파이썬을 파이썬 답게 강의를 기반으로, 문제를 풀며 알게 된 함수들을 추가로 정리하였다.

기초

iterable, sequence

iterable

자신의 멤버를 한 번에 하나씩 리턴할 수 있는 객체

• str
• list
• tuple
• dict

sequence

int 타입 인덱스를 통해 원소에 접근할 수 있는 iterable

• str
• list
• tuple

* (unpacking 연산자)

함수 인자(arguments)에 사용해서 list 등의 묶여있는 자료구조를 풀어서 전달하는 연산자

l = [1, 2, 3, 4]

print(l)	# [1, 2, 3, 4]
print(*l)	# 1 2 3 4

문자열 공백 포함 정렬

s = '가나'

# 왼쪽 정렬
print(s.ljust(5))	# '가나   '

# 오른쪽 정렬
print(s.rjust(5))	# '   가나'

# 가운데 정렬
print(s.center(5))	# '  가나 '

알파벳, 숫자 모두

import string

# 알파벳 lowercase
print(string.ascii_lowercase)	# 'abcdefghijklmnopqrstuvwxyz'

# 알파벳 uppercase
print(string.ascii_uppercase)	# 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

# 알파벳 전부
print(string.ascii_letters)	# 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'

# 숫자
print(string.digits)	# '0123456789'

collections

deque

stack 또는 queue가 필요할 때 유용

from collections import deque

# stack 처럼 사용
stack = deque()

stack.append(1)
stack.append(2)

print(list(stack))	# [1, 2]

stack.pop()

print(list(stack))	# [1]

# queue 처럼 사용
queue = deque()

queue.append(1)
queue.append(2)

print(list(queue))	# [1, 2]

queue.popleft()

print(list(queue))	# [2]

Counter

dict를 상속받는 클래스로, 개수를 셀 때 유용

from collections import Counter

mystr = 'aabbc'

d = Counter(mystr)
print(d)	# Counter({'a': 2, 'b': 2, 'c': 1})

# 출력 결과 : a b c
for key in d.keys():
	print(key, end=' ')
print()

# 출력 결과 : 2 2 1
for value in d.values():
	print(value, end=' ')
print()

# 출력 결과 : a: 2 b: 2 c: 1
for key, value in d.items():
	print(f'{key}: {value}', end=' ')
print()

heapq

priority queue 또는 heap이 필요할 때 유용
기본은 작은 값이 우선순위가 높은 최소힙(min-heap)

import heapq

l = []

# 힙 방식으로 삽입. 첫 번째 인자 list, 두 번째 인자 삽입할 item
heapq.heappush(l, 1)
heapq.heappush(l, 4)
heapq.heappush(l, 2)

# 힙 방식으로 삭제. heappop할 list를 인자로 전달
print(heapq.heappop(l))	# 1
print(heapq.heappop(l))	# 2
print(heapq.heappop(l))	# 4

h = [5, 2, 3, 6, 1]

# 힙이 아닌 list를 힙으로 만들기. 힙으로 만들 list를 인자로 전달
heapq.heapify(h)

print(heapq.heappop(h))	# 1
print(heapq.heappop(h))	# 2
print(heapq.heappop(h))	# 3
print(heapq.heappop(h))	# 5
print(heapq.heappop(h))	# 6

itertools

chain(*iterables)

여러 개의 iterable을 이어서 하나의 iterable로 만드는 함수

from itertools import chain

print(list(chain([1, 2, 3, 4], [5, 6, 7, 8])))	# [1, 2, 3, 4, 5, 6, 7, 8]

product

iterable의 product(데카르트 곱)를 구하는 함수

from itertools import product

A = 'ABCD'
x = 'xy'

print(list(product(A, x)))
# [('A', 'x'), ('A', 'y'), ('B', 'x'), ('B', 'y'), ('C', 'x'), ('C', 'y'), ('D', 'x'), ('D', 'y')]

permutations

iterable의 permutation(순열)을 구하는 함수

from itertools import permutations

print(list(permutations(['1', '2', '3'])))
# [('1', '2', '3'), ('1', '3', '2'), ('2', '1', '3'), ('2', '3', '1'), ('3', '1', '2'), ('3', '2', '1')]

combinations

iterable의 combination(조합)을 구하는 함수
첫 번째 인자로 iterable, 두 번째 인자로 뽑을 개수 전달

from itertools import combinations

print(list(cominations(['1', '2', '3'], 2)))
# [('1', '2'), ('1', '3'), ('2', '3')]

내장 함수

divmod(a, b)

a를 b로 나눈 몫과 나머지를 반환하는 함수
단, 작은 숫자는 그냥 a//b, a%b를 쓰는 게 더 빠르다고 한다.

print(divmod(5, 3))	# (1, 2)

any(iterable) / all(iterable)

any는 iterable 객체를 받아서 그 객체를 돌며 True가 하나라도 있으면 True를 반환하는 함수이다.

print(any([True, False, False]))	# True
print(any([False, False, False]))	# False

all은 iterable 객체를 받아서 그 객체의 모든 요소가 True일 때만 True를 반환하는 함수이다.

print(all([True, True, True]))	# True
print(any([True, True, False]))	# False

zip(*iterables)

zip은 여러 개의 iterable을 인자로 받아서 각 iterable의 iterator 묶음을 가져오는 함수이다. 여러 iterable을 동시에 순회할 때 사용하면 유용하다.
zip에 인자로 넘겨지는 iterable의 길이가 다르면 짧은 것을 기준으로 순회 후 종료한다.

• 예) zip 함수로 한 리스트의 인접한 원소간의 차 구하기

answer = []
mylist = [2, 6, 4, 8, 1]
for n1, n2 in zip(mylist, mylist[1:]):
	answer.append(abs(n1-n2))

print(answer)	# [4, 2, 4, 7]

transpos 구현 방법

mylist = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print(list(map(list, zip(*mylist)))	# [[1, 4, 7], [2, 5, 8], [3, 6, 9]]

# 원리 이해
print(list(zip(*mylist)))	# [(1, 4, 7), (2, 5, 8), (3, 6, 9)]

2차원 list를 1차원 list로 바꾸는 방법

import itertools

mylist = [[1, 2], [3, 4]]

# 1
print(sum(mylist, []))	# [1, 2, 3, 4]

# 2
print(list(itertools.chain(*mylist))	# [1, 2, 3, 4]

정렬

기본적인 정렬 방법은 list.sort() 함수를 사용하는 방법, sorted() 내장함수를 사용하는 방법 2가지가 있다.

• sort() 함수는 list에만 있는 메서드다.
  • sort() 함수는 원본 list 자체를 정렬된 상태로 변경시킨다.
• sorted() 함수는 iterable이면 모두 정렬 가능하다.
  • sorted() 함수는 iterable을 정렬한 list를 반환한다.
• 둘 다 기본적으로 오름차순 정렬을 한다.

l = [-1, 5, 0, 2, 4]
l.sort()
print(l)	# [-1, 0, 2, 4, 5]

l2 = [-4, 0, 6, 8, 2]
print(sorted(l2))

특정 항목에 대한 정렬

l = [(0, 2), (5, -4), (-1, 2), (0, 1), (7, 1)]

# 기본적으로 앞의 element -> 뒤의 element의 오름차순으로 정렬
print(sorted(l))	# [(-1, 2), (0, 1), (0, 2), (5, -4), (7, 1)]

# 2번째 element를 기준으로 정렬
print(sorted(l, key=lambda x: x[1]))	# [(5, -4), (0, 1), (7, 1), (0, 2), (-1, 2)]

• key에 전달하는 함수는 인자를 한 개 가지고 정렬의 기준이 될 값을 반환

정렬 기준 여러개 사용하기

from operator import itemgetter
l = [(0, 2), (5, -4), (-1, 2), (0, 1), (7, 1)]

# 1번 element를 기준으로 정렬하고, 1번 element가 같으면 0번 element를 기준으로 정렬
print(sorted(l, key=itemgetter(1, 0)))	# [(5, -4), (0, 1), (7, 1), (-1, 2), (0, 2)]

장식-정렬-복원(Decorate-Sort-Undecorate)

1. 정렬 순서 상 앞에 와야 하는 것으로 장식(decorate)한 새로운 배열을 만든다.
2. 장식된 배열을 정렬한다.
3. 장식된 배열에서 필요한 부분만 추출해 새로운 배열로 만든다. 이 배열은 원하는 순서대로 정렬되어 있다.

• 단어의 길이 순으로 정렬하려고 하는 등 좀 더 복잡한 방식의 정렬을 할 때 유용한 방법이다.

words = [ 'i', 'am', 'a', 'Student' ]

# decorate
decorated = [ (len(word), word) for word in words ]
# sort
decorated.sort()
# undecorate
sorted_words = [ word for _, word in decorated ]

print(sorted_words)	# ['a', 'i', 'am', 'Student']