(코딜리티) Lesson 5 : PassingCars

문제 : app.codility.com/programmers/lessons/5-prefix_sums/passing_cars/

PassingCars coding task - Learn to Code - Codility

문제 해설

비어있지 않은 N개인 배열 A가 주어 질 때, 0과 1인 경우의 쌍 갯수를 구하시오.

 

예를 들어, A = [0, 1, 0 ,1, 1]일 경우,

0에서 1인 쌍은 (0,1), (0, 3), (0, 4), (2, 3), (2, 4)로 5를 return해야 합니다.

그리고 만약 총 갯수가 1,000,000,000개를 초과할 경우는 -1을 return합니다.

제한 사항

N은 1~100,000사이의 값

A의 각 요소는 0또는 1

코드

def solution(A):
    
    prefix_sum = [0]

    for a in A:
        b = prefix_sum[-1]
        b += a
        prefix_sum.append(b)

    result = 0
    for idx, i in enumerate(A):
        if i == 0:
            r = prefix_sum[-1] - prefix_sum[idx+1]
            result += r

    return result if result <=1000000000 else -1

해설

prefix sum 문제로 배열 A를 순서대로 뽑아보면서 0인경우의 index번호에서 1의 갯수만 세아려서 계속 더해주면 됩니다.