20_bitmask_dp.py - Examples

  1"""
  2비트마스크 DP (Bitmask Dynamic Programming)
  3Bitmask DP
  4
  5비트 연산을 활용하여 집합 상태를 표현하는 DP 기법입니다.
  6"""
  7
  8from typing import List, Tuple
  9from functools import lru_cache
 10
 11
 12# =============================================================================
 13# 1. 비트 연산 기초
 14# =============================================================================
 15
 16def bit_operations_demo():
 17    """비트 연산 기본"""
 18    n = 5  # 집합 크기
 19
 20    # 빈 집합
 21    empty = 0
 22
 23    # 전체 집합 {0, 1, 2, 3, 4}
 24    full = (1 << n) - 1  # 11111 (이진수)
 25
 26    # i번째 원소 추가
 27    def add(mask: int, i: int) -> int:
 28        return mask | (1 << i)
 29
 30    # i번째 원소 제거
 31    def remove(mask: int, i: int) -> int:
 32        return mask & ~(1 << i)
 33
 34    # i번째 원소 토글
 35    def toggle(mask: int, i: int) -> int:
 36        return mask ^ (1 << i)
 37
 38    # i번째 원소 포함 여부
 39    def contains(mask: int, i: int) -> bool:
 40        return bool(mask & (1 << i))
 41
 42    # 원소 개수
 43    def count(mask: int) -> int:
 44        return bin(mask).count('1')
 45
 46    # 최하위 비트 (가장 작은 원소)
 47    def lowest_bit(mask: int) -> int:
 48        return mask & (-mask)
 49
 50    # 부분집합 순회
 51    def subsets(mask: int):
 52        """mask의 모든 부분집합을 순회"""
 53        subset = mask
 54        while True:
 55            yield subset
 56            if subset == 0:
 57                break
 58            subset = (subset - 1) & mask
 59
 60    return {
 61        'empty': empty,
 62        'full': full,
 63        'add': add,
 64        'remove': remove,
 65        'toggle': toggle,
 66        'contains': contains,
 67        'count': count,
 68        'lowest_bit': lowest_bit,
 69        'subsets': subsets
 70    }
 71
 72
 73# =============================================================================
 74# 2. 외판원 문제 (TSP - Traveling Salesman Problem)
 75# =============================================================================
 76
 77def tsp(dist: List[List[int]]) -> int:
 78    """
 79    외판원 문제 (TSP)
 80    모든 도시를 방문하고 시작점으로 돌아오는 최소 비용
 81
 82    시간복잡도: O(n² * 2^n)
 83    공간복잡도: O(n * 2^n)
 84    """
 85    n = len(dist)
 86    INF = float('inf')
 87
 88    # dp[mask][i] = mask에 포함된 도시를 방문하고 현재 i에 있을 때 최소 비용
 89    dp = [[INF] * n for _ in range(1 << n)]
 90    dp[1][0] = 0  # 시작점 (도시 0)
 91
 92    for mask in range(1 << n):
 93        for last in range(n):
 94            if dp[mask][last] == INF:
 95                continue
 96            if not (mask & (1 << last)):
 97                continue
 98
 99            for next_city in range(n):
100                if mask & (1 << next_city):
101                    continue
102
103                new_mask = mask | (1 << next_city)
104                dp[new_mask][next_city] = min(
105                    dp[new_mask][next_city],
106                    dp[mask][last] + dist[last][next_city]
107                )
108
109    # 모든 도시 방문 후 시작점으로
110    full_mask = (1 << n) - 1
111    result = min(dp[full_mask][i] + dist[i][0] for i in range(n))
112
113    return result if result != INF else -1
114
115
116def tsp_path(dist: List[List[int]]) -> Tuple[int, List[int]]:
117    """TSP 최소 비용과 경로 반환"""
118    n = len(dist)
119    INF = float('inf')
120
121    dp = [[INF] * n for _ in range(1 << n)]
122    parent = [[-1] * n for _ in range(1 << n)]
123
124    dp[1][0] = 0
125
126    for mask in range(1 << n):
127        for last in range(n):
128            if dp[mask][last] == INF:
129                continue
130
131            for next_city in range(n):
132                if mask & (1 << next_city):
133                    continue
134
135                new_mask = mask | (1 << next_city)
136                new_cost = dp[mask][last] + dist[last][next_city]
137
138                if new_cost < dp[new_mask][next_city]:
139                    dp[new_mask][next_city] = new_cost
140                    parent[new_mask][next_city] = last
141
142    full_mask = (1 << n) - 1
143    min_cost = INF
144    last_city = -1
145
146    for i in range(n):
147        cost = dp[full_mask][i] + dist[i][0]
148        if cost < min_cost:
149            min_cost = cost
150            last_city = i
151
152    # 경로 복원
153    path = []
154    mask = full_mask
155    city = last_city
156
157    while city != -1:
158        path.append(city)
159        prev_city = parent[mask][city]
160        mask ^= (1 << city)
161        city = prev_city
162
163    path.reverse()
164    path.append(0)  # 시작점으로 복귀
165
166    return min_cost, path
167
168
169# =============================================================================
170# 3. 집합 분할 문제 (Set Partition)
171# =============================================================================
172
173def can_partition_k_subsets(nums: List[int], k: int) -> bool:
174    """
175    배열을 합이 같은 k개의 부분집합으로 분할 가능한지
176    시간복잡도: O(n * 2^n)
177    """
178    total = sum(nums)
179    if total % k != 0:
180        return False
181
182    target = total // k
183    n = len(nums)
184
185    # dp[mask] = mask 집합을 사용했을 때 현재 버킷의 합 (target으로 나눈 나머지)
186    dp = [-1] * (1 << n)
187    dp[0] = 0
188
189    for mask in range(1 << n):
190        if dp[mask] == -1:
191            continue
192
193        for i in range(n):
194            if mask & (1 << i):
195                continue
196
197            if dp[mask] + nums[i] <= target:
198                new_mask = mask | (1 << i)
199                dp[new_mask] = (dp[mask] + nums[i]) % target
200
201    return dp[(1 << n) - 1] == 0
202
203
204# =============================================================================
205# 4. 최소 비용 작업 할당 (Assignment Problem)
206# =============================================================================
207
208def min_cost_assignment(cost: List[List[int]]) -> int:
209    """
210    n명의 사람에게 n개의 작업을 1:1 할당하는 최소 비용
211    cost[i][j] = 사람 i가 작업 j를 수행하는 비용
212
213    시간복잡도: O(n * 2^n)
214    """
215    n = len(cost)
216
217    @lru_cache(maxsize=None)
218    def dp(mask: int) -> int:
219        person = bin(mask).count('1')
220
221        if person == n:
222            return 0
223
224        min_cost = float('inf')
225        for job in range(n):
226            if mask & (1 << job):
227                continue
228
229            min_cost = min(min_cost, cost[person][job] + dp(mask | (1 << job)))
230
231        return min_cost
232
233    return dp(0)
234
235
236# =============================================================================
237# 5. 해밀턴 경로 (Hamiltonian Path)
238# =============================================================================
239
240def hamiltonian_path_count(adj: List[List[int]]) -> int:
241    """
242    해밀턴 경로의 개수 (모든 정점을 한 번씩 방문하는 경로)
243    adj: 인접 행렬 (adj[i][j] = 1이면 i→j 간선 존재)
244
245    시간복잡도: O(n² * 2^n)
246    """
247    n = len(adj)
248
249    # dp[mask][i] = mask 정점들을 방문하고 i에서 끝나는 경로 수
250    dp = [[0] * n for _ in range(1 << n)]
251
252    # 초기화: 각 정점에서 시작
253    for i in range(n):
254        dp[1 << i][i] = 1
255
256    for mask in range(1 << n):
257        for last in range(n):
258            if dp[mask][last] == 0:
259                continue
260            if not (mask & (1 << last)):
261                continue
262
263            for next_v in range(n):
264                if mask & (1 << next_v):
265                    continue
266                if not adj[last][next_v]:
267                    continue
268
269                new_mask = mask | (1 << next_v)
270                dp[new_mask][next_v] += dp[mask][last]
271
272    # 모든 정점 방문한 경로 합
273    full_mask = (1 << n) - 1
274    return sum(dp[full_mask])
275
276
277# =============================================================================
278# 6. 스티커 최적 배치 (SOS DP 전처리)
279# =============================================================================
280
281def sos_dp(arr: List[int]) -> List[int]:
282    """
283    Sum over Subsets DP
284    각 마스크에 대해 부분집합들의 값 합 계산
285
286    result[mask] = sum(arr[subset]) for all subset of mask
287
288    시간복잡도: O(n * 2^n)
289    """
290    n = len(arr).bit_length()
291    dp = arr.copy()
292
293    # 0~(len(arr)-1)까지 확장
294    while len(dp) < (1 << n):
295        dp.append(0)
296
297    for i in range(n):
298        for mask in range(1 << n):
299            if mask & (1 << i):
300                dp[mask] += dp[mask ^ (1 << i)]
301
302    return dp
303
304
305# =============================================================================
306# 7. 최대 독립 집합 (Maximum Independent Set on Trees - Bitmask)
307# =============================================================================
308
309def max_independent_set(adj: List[List[int]]) -> int:
310    """
311    그래프에서 최대 독립 집합 크기 (서로 인접하지 않은 정점 집합)
312    작은 그래프에서 비트마스크로 brute force
313
314    시간복잡도: O(2^n * n²)
315    """
316    n = len(adj)
317    max_size = 0
318
319    for mask in range(1 << n):
320        # mask가 독립 집합인지 확인
321        valid = True
322        for i in range(n):
323            if not (mask & (1 << i)):
324                continue
325            for j in range(i + 1, n):
326                if not (mask & (1 << j)):
327                    continue
328                if adj[i][j]:
329                    valid = False
330                    break
331            if not valid:
332                break
333
334        if valid:
335            max_size = max(max_size, bin(mask).count('1'))
336
337    return max_size
338
339
340# =============================================================================
341# 8. 격자 채우기 (Broken Profile DP)
342# =============================================================================
343
344def domino_tiling(m: int, n: int) -> int:
345    """
346    m×n 격자를 1×2 도미노로 채우는 경우의 수
347    비트마스크 DP (profile 방식)
348
349    시간복잡도: O(n * 2^m * 2^m)
350    """
351    if m > n:
352        m, n = n, m
353
354    # dp[col][profile] = 현재 열까지 채우고 프로파일이 profile인 경우의 수
355    dp = {0: 1}
356
357    for col in range(n):
358        for row in range(m):
359            new_dp = {}
360
361            for profile, count in dp.items():
362                # 현재 셀이 이미 채워진 경우
363                if profile & (1 << row):
364                    new_profile = profile ^ (1 << row)
365                    new_dp[new_profile] = new_dp.get(new_profile, 0) + count
366                else:
367                    # 수평 도미노 (다음 열로 확장)
368                    new_profile = profile | (1 << row)
369                    new_dp[new_profile] = new_dp.get(new_profile, 0) + count
370
371                    # 수직 도미노 (아래 셀과 함께)
372                    if row + 1 < m and not (profile & (1 << (row + 1))):
373                        new_dp[profile] = new_dp.get(profile, 0) + count
374
375            dp = new_dp
376
377    return dp.get(0, 0)
378
379
380# =============================================================================
381# 테스트
382# =============================================================================
383
384def main():
385    print("=" * 60)
386    print("비트마스크 DP (Bitmask DP) 예제")
387    print("=" * 60)
388
389    # 1. 비트 연산 기초
390    print("\n[1] 비트 연산 기초")
391    ops = bit_operations_demo()
392    mask = 0b10110  # {1, 2, 4}
393    print(f"    mask = {bin(mask)} ({mask})")
394    print(f"    원소 개수: {ops['count'](mask)}")
395    print(f"    3 포함: {ops['contains'](mask, 3)}")
396    print(f"    2 포함: {ops['contains'](mask, 2)}")
397    print(f"    3 추가: {bin(ops['add'](mask, 3))}")
398    print(f"    부분집합: ", end="")
399    for s in ops['subsets'](mask):
400        print(f"{bin(s)} ", end="")
401    print()
402
403    # 2. TSP
404    print("\n[2] 외판원 문제 (TSP)")
405    dist = [
406        [0, 10, 15, 20],
407        [10, 0, 35, 25],
408        [15, 35, 0, 30],
409        [20, 25, 30, 0]
410    ]
411    min_cost, path = tsp_path(dist)
412    print(f"    거리 행렬: 4x4")
413    print(f"    최소 비용: {min_cost}")
414    print(f"    경로: {path}")
415
416    # 3. 집합 분할
417    print("\n[3] K개 부분집합 분할")
418    nums = [4, 3, 2, 3, 5, 2, 1]
419    k = 4
420    result = can_partition_k_subsets(nums, k)
421    print(f"    배열: {nums}, k={k}")
422    print(f"    분할 가능: {result}")
423
424    # 4. 작업 할당
425    print("\n[4] 최소 비용 작업 할당")
426    cost = [
427        [9, 2, 7, 8],
428        [6, 4, 3, 7],
429        [5, 8, 1, 8],
430        [7, 6, 9, 4]
431    ]
432    min_assign = min_cost_assignment(cost)
433    print(f"    비용 행렬: 4x4")
434    print(f"    최소 비용: {min_assign}")
435
436    # 5. 해밀턴 경로
437    print("\n[5] 해밀턴 경로 개수")
438    adj = [
439        [0, 1, 1, 1],
440        [1, 0, 1, 0],
441        [1, 1, 0, 1],
442        [1, 0, 1, 0]
443    ]
444    count = hamiltonian_path_count(adj)
445    print(f"    인접 행렬: 4x4")
446    print(f"    해밀턴 경로 수: {count}")
447
448    # 6. 최대 독립 집합
449    print("\n[6] 최대 독립 집합")
450    adj2 = [
451        [0, 1, 0, 1],
452        [1, 0, 1, 0],
453        [0, 1, 0, 1],
454        [1, 0, 1, 0]
455    ]
456    mis = max_independent_set(adj2)
457    print(f"    그래프: 4-cycle")
458    print(f"    최대 독립 집합 크기: {mis}")
459
460    # 7. 도미노 타일링
461    print("\n[7] 도미노 타일링")
462    for m, n in [(2, 3), (2, 4), (3, 4)]:
463        count = domino_tiling(m, n)
464        print(f"    {m}×{n} 격자: {count}가지")
465
466    # 8. SOS DP
467    print("\n[8] SOS DP (부분집합 합)")
468    arr = [1, 2, 4, 8]  # 각 원소는 해당 비트의 값
469    result = sos_dp(arr)
470    print(f"    배열: {arr}")
471    print(f"    result[0b0111] = result[7] = {result[7]}")
472    print(f"    (부분집합: {{0},{1},{0,1},{2},{0,2},{1,2},{0,1,2}} 합)")
473
474    print("\n" + "=" * 60)
475
476
477if __name__ == "__main__":
478    main()