Codility: PassingCars (Python)

주제

Lesson 5: Prefix Sums

 

난이도

Easy

 

문제

A non-empty array A consisting of N integers is given. The consecutive elements of array A represent consecutive cars on a road.
Array A contains only 0s and/or 1s:
0 represents a car traveling east,1 represents a car traveling west.
The goal is to count passing cars. We say that a pair of cars (P, Q), where 0 ≤ P < Q < N, is passing when P is traveling to the east and Q is traveling to the west.

For example, consider array A such that:
A[0] = 0
A[1] = 1
A[2] = 0
A[3] = 1
A[4] = 1

We have five pairs of passing cars: (0, 1), (0, 3), (0, 4), (2, 3), (2, 4).
Write a function:
def solution(A)
that, given a non-empty array A of N integers, returns the number of pairs of passing cars.
The function should return −1 if the number of pairs of passing cars exceeds 1,000,000,000.

For example, given:
A[0] = 0
A[1] = 1
A[2] = 0
A[3] = 1
A[4] = 1

the function should return 5, as explained above.

Write an efficient algorithm for the following assumptions:
 - N is an integer within the range [1..100,000];
 - each element of array A is an integer that can have one of the following values: 0, 1.

 

서로 마주치며 지나가는 차의 쌍(pair)을 세는 문제이다. 참고로 0은 동쪽, 1은 서쪽으로 향하는 차량을 의미한다. 

여기서 N은 배열의 길이, 즉 차량의 총 수를 의미하고, 1 ~ 100,000대 가능하다. 또한 배열의 원소는 0 또는 1 뿐이다. 

또한 efficient해야한다.

 

풀이

역시 efficient하게 풀라는 문제는 inefficient하게 푼 다음에 효율적으로 바꿔야 제맛이다! (아니다. 바로 효율적으로 풀수 있는 두뇌가 제일 부럽다.) 따라서 O(N^2)과 O(N) 풀이 둘 다 준비했다.

 

1) O(N^2)

머리를 쓴다고 열심히 쥐어짜낸 생각은 아래와 같다.

 - 0과 1이 만나면 passing이며, 순차적으로 생각해야한다. (ex. A[1] = 1인 차량과 A[2] = 0 인 차량은 서로 passing으로 보지 않음) 

 - 그러므로 앞에서부터 0을 만나면 그 뒤에 1의 갯수를 세서 계속 합하자.

 

코드가 매우 간단하게 짜진다. 문제는 python의 count함수는 O(N)이므로, 배열 A의 원소를 계속 반복해서 0인지 아닌지 판단하는 과정 속에 count까지 더해지면서 O(N^2)이 된다는 것이다. 즉, 비효율적이다.

 

2) O(N)

prefix sum 단원이기 때문에 prefix를 사용해야할 것 같지만 실제로 사용하진 않았다. (아닌가 사용된건가? 아닌 것 같다.)

그냥 열심히 예시를 바라보면서 반복문 하나로 어떻게 처리해야할까 고민하다가 문득 떠올랐다. 

 

 - 0이 늘어날 때 마다, 수를 하나씩 늘려 passing_car에 차곡차곡 쌓자.

 

이게 무슨 말이냐면, 아래와 같다. 

 

 - 0번 차량 : 0(east) 

 - 1번 차량 : 1(west). 이때 (0번 차량,1번 차량)으로 한 번 passing한다. (total passing : 1)

 - 2번 차량 : 0. 즉, 현재 east행이 2 대 있다.

 - 3번 차량 : 1. 이때, 3번 차량과 passing하는 east 행 차량은 2대이다. (total passing : 1 + 2)

 - 4번 차량 : 1. 4번 차량과 passing하는 east 행 차량은 2대이다. (total passing : 1 + 2 + 2)

 

즉, 위처럼 0의 갯수가 늘어나면 1을 만날 때 마다 0의 갯수 만큼 계속 합을 한다는 것이다.  

 

차량 번호	차량 방향	차량과 passing하는 순서쌍	차량과 passing하는 차량 수	총 passing cars
0	0	-	0	0
1	1	(0,1)	1	1
2	0	-	0	1
3	1	(0,3), (2,3)	2	3
4	1	(0,4), (2,4)	2	5

 

 

코드

1) O(N^2)

def solution_n2(A):
    total_passing = 0
    # 0뒤로 배열을 나눠서, 뒷배열에 1이 발생하는 숫자를 카운트해서 다 합하는 방식
    # count함수가 O(N)이므로 결국 n^2이 된다
    for idx, pos in enumerate(A):
        if pos == 0:
            passing = A[idx:].count(1)
            total_passing += passing

    if total_passing > 1000000000:
        total_passing = -1

    return total_passing

 

2) O(N)

def solution(A):
    zero_count = 0
    passing_car = 0

    # 반복문 한번만 돌아야함
    # 그래서 차근차근 예시를 보면, 맨 처음에 0이고, 뒤에 1이므로, (0,1)로 한번 passing
    # 그 후 다시 0이 나오므로, east행이 2개임.
    # 그후에 1이 아오는 데, 이때 0번 차량과 2번 차량이 3번 차량을 passing하므로 +2 해줘야함.
    # 맨 마지막도 1인데, 이것도 0번과 2번 차량을 passing하므로 +2
    # 즉, 0의 갯수가 늘면 1을 만날때마다 0의 갯수만큼 +해주면 된다.

    for i in A:
        if i == 0:
            zero_count +=1
        else:
            passing_car += zero_count

    if passing_car > 1000000000:
        passing_car = -1
    return passing_car

 

시간복잡도

길이가 N인 배열 A를 한바퀴 반복하는 동안 모든 연산이 끝난다.

따라서 O(N)의 시간복잡도를 갖는다.

 

 

결과

Codility 결과

 

이틀 전에 풀고 이제서야 풀이과정을 올리는데, 이 당시 내 두뇌는 좀 잘돌아갔던 것 같다...ㅋㅋㅋㅋㅋㅋㅋ 

제발 감을 잃지 않도록..!