[Programmers] #181931 - 등차수열의 특정한 항만 더하기 [Java][C++][Python]

문제 링크

---

1. 아이디어

첫째항 a, 공차 d인 등차수열의 $(i+1)$번째 항은 $a + d \cdot i$($i$는 0-indexed)로 바로 구할 수 있다. included 배열을 순회하면서 included[i]가 true인 위치의 항만 이 식으로 계산해 더하면 된다.

---

2. 복잡도

접근	시간	공간
풀이	$O(N)$	$O(1)$

($N$ = included의 길이)

---

3. 코드

풀이 [Java][C++][Python]

class Solution {
    public int solution(int a, int d, boolean[] included) {
        int sum = 0;
        for (int i = 0; i < included.length; i++) {
            if (included[i]) sum += a + d * i;
        }

        return sum;
    }
}

#include <bits/stdc++.h>
using namespace std;

int solution(int a, int d, vector<bool> included) {
    int sum = 0;
    for (int i = 0; i < included.size(); i++) {
        if (included[i]) sum += a + d * i;
    }

    return sum;
}

def solution(a, d, included):
    return sum(a + d * i for i, flag in enumerate(included) if flag)

---