[leetcode #413] Arithmetic Slices

Problem

An integer array is called arithmetic if it consists of at least three elements and if the difference between any two consecutive elements is the same.

・ For example, [1,3,5,7,9], [7,7,7,7], and [3,-1,-5,-9] are arithmetic sequences.

Given an integer array nums, return the number of arithmetic subarrays of nums.

A subarray is a contiguous subsequence of the array.

Example 1:

Input: nums = [1,2,3,4]
Output: 3
Explanation: We have 3 arithmetic slices in nums: [1, 2, 3], [2, 3, 4] and [1,2,3,4] itself.

Example 2:

Input: nums = [1]
Output: 0

Constraints:

・ 1 <= nums.length <= 5000
・ -1000 <= nums[i] <= 1000

Idea

주어진 배열에서 arithmetic인 subarray의 갯수를 구하는 문제다.

연속된 두 수의 차이가 계속해서 똑같을 경우 연속된 수의 개수를 1씩 올린다. 이 개수를 n이라고 하자.

만약 차가 다른 구간이 나오고 n이 3보다 크거나 같으면 총합에 (n-2) * (n-1) / 2를 더하면 된다. 연속된 수가 5라고 할 때, arithmetic인 subarray는 다음과 같다.

[1,2,3]
[2,3,4]
[3,4,5]
[1,2,3,4]
[2,3,4,5]
[1,2,3,4,5]

n이 5일 때, 3+2+1 = (3+1) 3 / 2 = (5-1) (5-2) / 2가 subarray의 갯수인 것이다.

차이가 다를 때 두 수의 차를 바뀐 값으로 대체하고 n은 2로 정한다.

루프를 빠져나온 뒤 n이 3보다 크거나 같으면 총합에 (n-2) * (n-1) / 2를 더한다.

Solution

class Solution {
    public int numberOfArithmeticSlices(int[] nums) {
        if (nums.length < 3) {
            return 0;
        }

        int diff = nums[1] - nums[0];
        int n = 2;
        int res = 0;

        for (int i=2; i < nums.length; i++) {
            if (diff != nums[i] - nums[i-1]) {
                if (n >= 3) {
                    res += (n-2) * (n-1) / 2;
                }
                n = 2;
                diff = nums[i] - nums[i-1];
                continue;
            }
            n++;
        }

        if (n >= 3) {
            res += (n-2) * (n-1) / 2;
        }

        return res;
    }
}

Reference

https://leetcode.com/problems/arithmetic-slices/