파이썬으로 네 개 리스트에서 목표값 이하 합을 만드는 고유 인덱스 조합 개수 구하기

숫자로 이루어진 네 개의 리스트 A, B, C, D와 하나의 목표값(target)이 주어졌다고 가정해 봅시다. 이때 A[i] + B[j] + C[k] + D[l] ≤ target을 만족하는 서로 다른 고유 인덱스 조합 i, j, k, l의 개수를 구하는 것이 문제입니다.

예를 들어 입력이 A = [3, 2], B = [5, 3], C = [1], D = [2, 3], target = 9라고 하면 출력은 3이 됩니다. 가능한 조합은 다음과 같습니다.

[3, 3, 1, 2]
[3, 3, 1, 3]
[2, 3, 1, 3]

문제 해결 접근 방식

이 문제는 단순히 네 겹의 반복문으로 모든 경우를 확인하면 시간 복잡도가 O(n⁴)가 되어 비효율적입니다. 대신 다음과 같은 전략을 사용하면 효율적으로 해결할 수 있습니다.

먼저 빈 임시 리스트(temp_list)를 준비합니다.
A와 B의 모든 요소 쌍에 대해 두 값의 합(A[i] + B[j])을 임시 리스트에 추가합니다.
임시 리스트를 오름차순으로 정렬합니다.
C와 D의 모든 요소 쌍에 대해 다음을 반복합니다.
- sum_cd := C[i] + D[j]
- sum_ab := target − sum_cd
- 임시 리스트에서 sum_ab 이하인 원소의 개수를 정답에 더합니다.
최종 정답(ans)을 반환합니다.

A와 B의 합들을 미리 정렬해 둔 후, C와 D의 각 조합에 대해 이진 탐색(bisect)을 활용하면 조건을 만족하는 개수를 빠르게 찾을 수 있습니다. 이렇게 하면 전체 시간 복잡도를 O(n² log n) 수준으로 낮출 수 있습니다.

구현 예제

from bisect import bisect_right

class Solution:
    def solve(self, A, B, C, D, target):
        temp_list = []
        for i in range(len(A)):
            for j in range(len(B)):
                temp_list.append(A[i] + B[j])

        temp_list.sort()

        ans = 0
        for i in range(len(C)):
            for j in range(len(D)):
                sum_cd = C[i] + D[j]
                sum_ab = target - sum_cd

                ans += bisect_right(temp_list, sum_ab)

        return ans

ob = Solution()
A = [3, 2]
B = [5, 3]
C = [1]
D = [2, 3]
target = 9
print(ob.solve(A, B, C, D, target))

입력

[3, 2], [5, 3], [1], [2, 3], 9