Post

[프로그래머스] 제곱 개수 배열

프로그래머스 '제곱 개수 배열' 문제 풀이

Posted
By bienew22
views 9 min read
[프로그래머스] 제곱 개수 배열

문제

코딩테스트 연습 - 제곱 개수 배열
알고리즘 문제 연습 카카오톡 친구해요! 프로그래머스 교육 카카오 채널을 만들었어요. 여기를 눌러, 친구 추가를 해주세요. 신규 교육 과정 소식은 물론 다양한 이벤트 소식을 가장 먼저 알려드립니다.
school.programmers.co.kr

문제 정리

요약
  • 길이가 N인 정수 배열 arr이 주어집니다.
  • 다음 규칙에 따라서 brr을 생성합니다.
    • arr의 인덱스 순서대로 arr[i]를 brr에 연속으로 arr[i]개씩 추가합니다.
  • 예를 들면 arr = [1, 3, 2]이면 brr = [1, 3, 3, 3, 2, 2]이 됩니다.
  • brr의 부분 배열의 양 끝을 나타내는 l과 r이 주어집니다.
  • 여기서 K와 C를 구하려고 합니다.
    • K : 해당 부분 배열의 모든 원소의 합.
    • C : 해당 부분 배열과 길이가 같으면서 k 값이 같은 부분 배열의 개수.
입력 조건
  • N : 1 ~ 105
  • arr[i] : 1 ~ 105
  • l, r : 1 ~ sum(arr)
  • sum(brr) < 1015
출력
[K, C] 반환.

문제 풀이

문제 자체를 보면 간단한 슬라이딩 원도우 문제입니다. 하지만 최악의 경우를 고려하면 arr의 길이가 100,000이고 모든 원소가 100,000이며, l = r = 1인 상황에서는 윈도우 하나씩 이동하면 시간 복잡도가 O(1010) = O(N2)로 시간 초과가 발생하게 됩니다.

따라서 좀 더 효율적으로 윈도우를 이동하여 찾는 방법이 필요합니다.

arr = [4, 1, 4], l = 1, r = 6인 경우 슬라이드 하는 과정을 보면 다음과 같습니다. 예시

사진을 보면 arr index 기준 l과 r의 인덱스가 변경하지 않는 경우 동일하게 값의 변경이 일어나는 것을 확인할 수 있습니다. 즉, 변화가 없는 구간을 한 번에 건너뛸 수 있다는 점을 이용하면 원도우을 빠르게 이동시킬 수 있습니다.

추가로 문제를 해결할 때 brr 배열을 직접 구성할 필요가 없습니다. 각 arr[i]는 brr 상에서 idx ~ idx + arr[i] 구간을 차지하고 있습니다. 여기서 idx = arr[0] + arr[1] ... + arr[i - 1] 이므로 누적합(prefix sum)으로 구할 수 있습니다.

풀이 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

def solution(arr, l, r):
    ans_value, ans_cnt = 0, 1

    arr_size = len(arr)

    # part1 : prefix sum으로 arr[i] 차지하는 brr index 구간 구하기.
    s_idx_list = [0, ]
    e_idx_list = [arr[0] - 1]

    for idx in range(1, arr_size):
        s_idx_list.append(e_idx_list[-1] + 1)
        e_idx_list.append(e_idx_list[-1] + arr[idx])

    # part2 : l과 r에 해당하는 arr의 인덱스 찾기
    base_s, base_e = -1, -1

    for idx in range(arr_size):
        if s_idx_list[idx] <= l - 1 <= e_idx_list[idx]:
            base_s = idx
        
        if s_idx_list[idx] <= r - 1 <= e_idx_list[idx]:
            base_e = idx

    # part3 : 현재 구간의 합 구하기.
    s, length, idx_s = l - 1, r - l + 1, base_s

    while length > 0:
        move = e_idx_list[idx_s] - s + 1

        if move > length:
            move = length

        ans_value += arr[idx_s] * move
        idx_s += 1
        s += move
        length -= move

    # part4 : 현재 구간에 오른쪽으로 슬라이딩
    s, e, offset = l - 1, r - 1, 0
    idx_s, idx_e = base_s, base_e

    while e < e_idx_list[-1]:
        # e에서 다음 base_e 까지 현재 남은 개수
        # = 한 번에 이동할 횟수
        move_cnt = e_idx_list[idx_e] - e

        # 이동 횟수가 0인 경우
        if move_cnt == 0:
            idx_e += 1
            move_cnt = 1

            # 끝에 도달한 경우.
            if idx_e == arr_size:
                break

        e += move_cnt

        # 이동 하기.
        while move_cnt > 0:
            # s에서 다음 base_e 까지 현재 남은 개수
            # = 변동 폭이 같은 구간
            now_move = e_idx_list[idx_s]  - s

            if now_move > move_cnt:
                now_move = move_cnt

            diff = arr[idx_e] - arr[idx_s]

            if now_move == 0:
                idx_s += 1
                now_move = 1

            new_offset = offset + diff * now_move

            # offset이 0에서 시작해서 0으로끝나는 경우.
            if offset == 0 and new_offset == 0:
                ans_cnt += now_move
            # offset * new_offset < 0 : 부호가 바뀐 경우, 0을 경유함을 의미, 
            # offset % diff == 0 : 기존 offset에 이동폭의 배수이면 0에 도착해서 이동했음을 의미.
            elif offset * new_offset < 0 and offset % diff == 0:
                ans_cnt += 1
            # 0에 도착한 경우.
            elif new_offset == 0:
                ans_cnt += 1

            # 값 갱신
            offset = new_offset
            s += now_move
            move_cnt -= now_move


    # part 5 : 현재 구간에 왼쪽으로 슬라이딩
    # 로직은 오른쪽으로 이동의 반대로 수행.
    s, e, offset = l - 1, r - 1, 0
    idx_s, idx_e = base_s, base_e

    while s > 0:
        move_cnt = s - s_idx_list[idx_s]

        if move_cnt == 0:
            idx_s -= 1
            move_cnt = 1

            if idx_s < 0:
                break

        s -= move_cnt

        while move_cnt > 0:
            now_move = e - s_idx_list[idx_e]

            if now_move > move_cnt:
                now_move = move_cnt

            diff = arr[idx_s] - arr[idx_e]

            if now_move == 0:
                idx_e -= 1
                now_move = 1

            new_offset = offset + diff * now_move

            if offset == 0 and new_offset == 0:
                ans_cnt += now_move
            elif offset * new_offset < 0 and offset % diff == 0:
                ans_cnt += 1
            elif new_offset == 0:
                ans_cnt += 1

            offset = new_offset
            e -= now_move
            move_cnt -= now_move

    return [ans_value, ans_cnt]

회고

현재 구간의 합을 구하는데 별 의심 없이 O(N)으로 해결할 수 있다고 생각하였습니다.

문제 코드
1
2
3
4
5
6
7

  s, e, idx_s = l - 1, r, base_s

  while s < e:
      if s > e_idx_list[idx_s]:
          idx_s += 1
      ans_value += arr[idx_s]
      s += 1
놓친 부분
최악의 경우 테스트 케이스 설계할 때 원도우의 이동량은 고려했지만, l = 1, r = 1010 인 경우를 생각하지 못하였습니다.

추가로

원도우를 묶어서 이동한 경우, 문제 풀이에서는 이동 방향의 끝으로 한 번에 이동하는 방식을 채용하였습니다. 다음 코드처럼 동일하게 값이 변경되는 구간만 묶어서 한 번에 처리하는 것이 코드가 더욱 간결합니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

# arg1: e에서 다음 base_e 까지 현재 남은 개수
# arg2: s에서 다음 base_e 까지 현재 남은 개수
move_cnt = min(e_idx_list[idx_e] - e, e_idx_list[idx_s] - s)

if move_cnt == 0 and left_e == 0:
    # 끝에 도착한 경우
    if idx_e + 1 == arr_size:
        break

    # idx_e 증가 시키야 하는 경우
    idx_e += 1

# 한 번 이동 시 증감량.
diff = arr[idx_e] - arr[idx_s]

if move_cnt == 0 and left_s == 0:
    idx_s += 1

if move_cnt == 0:
    move_cnt = 1

new_offset = offset + diff * move_cnt
Problem Solving, Programmers
programmers sliding window prefix sum
This post is licensed under CC BY 4.0 by the author.