[코테] 코테에 유용한 파이썬 조각 모음

 

※ 대체로 코딩테스트 문제의 시간 제한은 1~5초 정도이며, 시간 제한이 명시되어 있지 않다면 대략 5초로 생각하자.

※ 코테에서 주로 사용되는 시간 복잡도는 선형 시간 O(N), 로그 선형 시간 O(NlogN), 이차 시간 O(N²)이다.

※ 파이썬의 들여쓰기는 스페이스 바(Space Bar)를 4번 입력하여 작성하자! (많은 파이썬 커뮤니티에서 사실상의 표준)

 

# 리스트 컴프리헨션은 튜플, 딕셔너리에서도 가능하다.

 

# 리스트 컴프리헨션을 사용한 2차원 리스트 초기화

 2차원 리스트 초기화에서는 반드시 리스트 컴프리헨션을 사용하자! (편리함!!!)

# N X M 크기의 2차원 리스트 초기화

n = 3
m = 3
array = [[0] * m for _ in range(n)]
print(array)

[[0, 0, 0], [0, 0, 0], [0, 0, 0]]

 

# Python Style remove_all() 이용법 (feat. 리스트 컴프리헨션)

 파이썬 리스트에서 특정한 값의 원소를 모두 제거하는 스타일은 다음 예시가 있다.

a = [1, 2, 3, 4, 5, 5, 5]
remove_set = {3, 5}

# remove_set에 포함되지 않은 값만을 저장
result = [i for i in a if i not in remove_set]
print(result)

[1, 2, 4]

 

# 집합 자료형

 - 집합 자료형의 연산

a = set([1, 2, 3, 4, 5])
b = set([3, 4, 5, 6, 7])

print(a | b) # 합집합
print(a & b) # 교집합
print(a - b) # 차집합

{1, 2, 3, 4, 5, 6, 7}
{3, 4, 5}
{1, 2}

 

 - 집합 자료형 관련 함수

data = set([1, 2, 3])
print(data)

# 새로운 원소 추가
data.add(4)
print(data)

# 새로운 원소 여러 개 추가
data.update([5, 6])
print(data)

# 특정한 값을 갖는 원소 삭제
data.remove(3)
print(data)

{1, 2, 3}
{1, 2, 3, 4}
{1, 2, 3, 4, 5, 6}
{1, 2, 4, 5, 6}

 

# 함수에서 사용되는 global 키워드

 해당 함수에서 지역 변수를 만들지 않고, 함수 바깥에 선언된 변수를 바로 참조한다.

a = 0

def func():
    global a
    a += 1
    
for i in range(10):
    func()
    
print(a)

10

 

# 람다 표현식(Lambda Express)

 함수를 매우 간단하게 작성하여 적용할 수 있다.

def add(a, b):
    return a + b

# 일반적인 add() 메소드 사용
print(add(3, 7))

# 람다 표현식으로 구현한 add() 메소드
print((lambda a, b: a + b)(3, 7))

10
10

 

# 여러 개의 반환 값

 파이썬은 여러 개의 반환 값을 가질 수 있다.

def operator(a, b):
    return a + b, a - b, a * b, a / b

a, b, c, d = operator(4, 2)
print(a, b, c, d)

6 2 8 2

 위의 operator 함수가 4 개의 값을 packing 하여 반환하면, 반환한 값을 a, b, c, d 변수에 unpacking하여 각각 할당한다.

 

# 입력을 위한 전형적인 소스코드

# 데이터의 개수 입력
n = int(input())

# 각 데이터를 공백으로 구분하여 입력
data = list(map(int, input().split()))

# 공백을 기준으로 구분하여 적은 수의 데이터 입력 (3개의 입력을 받는 경우)
a, b, c = map(int, input().split())

 

# 2차원 리스트 입력 받기

'''
3
4
0 0 0 0
0 0 0 0
0 0 0 0
'''

n = int(input())
m = int(input())

arr = []
for i in range(n):
    arr.append(list(map(int, input().split()))

# 또 하나의 방법
n = int(input())
m = int(input())

arr = [list(map(int, input().split())) for _ in range(n)]

 

# 입력을 최대한 빠르게 받아야 하는 경우 (입력만으로 시간초과 되는 경우)

import sys
sys.stdin.readline().rstrip()

 

# 파이썬 print() 줄 바꿈 수행 안하기

 파이썬의 print() 함수는 기본적을 출력 이후 줄 바꿈을 수행한다. 줄 변경을 피하고 싶다면, 콤마(,)를 사용하거나 end 속성을 이용한다.

print(8, 10)
print(8, end=' ')
print(10)

8 10
8 10

 

# 문자열과 수를 함께 출력하기

 단순히 더하기 연산자(+)를 사용하면 오류가 발생하기에 str() 함수를 이용하거나 콤마(,)를 사용한다.

answer = 7
print("정답은 " + str(answer) + "입니다.")
print("정답은", answer, "입니다.") # 공백이 생김을 유의

정답은 7입니다.
정답은 7 입니다.

 

# eval() 내장 함수

 수학 수식이 문자열 형식으로 들어올 때, 해당 수식을 계산한 결과를 반환한다.

result = eval("(3 + 5) * 7")
print(result)

56

 

# sorted() 내장 함수의 key 속성을 이용해 특정 기준에 따라 정렬하기

 리스트 내 튜플의 두 번째 원소를 기준으로 내림차순 정렬 하고자 할 때, 다음과 같이 실행할 수 있다.

result = sorted([('홍길동', 35), ('이순신', 75), ('아무개', 50)], key = lambda x: x[1], reverse = True)
print(result)

[('이순신', 75), ('아무개', 50), ('홍길동', 35)]

 

# itertools 라이브러리

 파이썬에서 반복되는 데이터를 처리하는데 유용한 라이브러리이다. 주로 순열, 조합에 많이 이용된다.

 

 - permutations (순열)

from itertools import permutations

data = ['A', 'B', 'C']
result = list(permutations(data, 3)) # 모든 순열 구하기

print(result)

[('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]

 

 - combinations (조합)

from itertools import combinations

data = ['A', 'B', 'C']
result = list(combinations(data, 2)) # 모든 조합 구하기

print(result)

[('A', 'B'), ('A', 'C'), ('B', 'C')]

 

 - product (중복 순열)

from itertools import product

data = ['A', 'B', 'C']
result = list(product(data, repeat=2)) # 2개를 뽑는 모든 순열 구하기 (중복 허용)

print(result)

[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'B'), ('B', 'C'), ('C', 'A'), ('C', 'B'), ('C', 'C')]

 

- combinations_with_replacement (중복 조합)

from itertools import combinations_with_replacement

data = ['A', 'B', 'C']
result = list(combinations_with_replacement(data, 2)) # 2개를 뽑는 모든 조합 구하기(중복 허용)

print(result)

[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

 

본 포스팅은 '안경잡이 개발자' 나동빈 님의

저서'이것이 코딩테스트다'와 그 유튜브 강의를 공부하고 정리한 내용을 담고 있습니다.