Python 暴力代码大模板：复制后直接改

  1#!/usr/bin/env python3
  2
  3import sys
  4from bisect import bisect_left, bisect_right
  5from collections import Counter, defaultdict, deque
  6from functools import cache
  7from heapq import heapify, heappop, heappush
  8from itertools import (
  9    accumulate,
 10    combinations,
 11    combinations_with_replacement,
 12    pairwise,
 13    permutations,
 14    product,
 15)
 16from math import gcd, inf, isqrt, lcm
 17
 18
 19# ============================================================
 20# 0. Constants and input
 21# ============================================================
 22
 23INF = inf
 24input = sys.stdin.buffer.readline
 25
 26
 27def read_ints(readline=None):
 28    """Read one line of whitespace-separated integers."""
 29    if readline is None:
 30        readline = input
 31    return list(map(int, readline().split()))
 32
 33
 34def read_all_ints(read=None):
 35    """Read all remaining whitespace-separated integers."""
 36    if read is None:
 37        read = sys.stdin.buffer.read
 38    return list(map(int, read().split()))
 39
 40
 41# Common solve() input patterns:
 42# n = int(input())
 43# a = read_ints()
 44# x, y = read_ints()
 45# data = read_all_ints()
 46# print(*a)
 47
 48
 49# ============================================================
 50# 1. List construction quick reference
 51# ============================================================
 52
 53# zeros = [0] * n
 54# squares = [x * x for x in range(n)]
 55# selected = [a[i] for i in range(n) if mask >> i & 1]
 56# grid = [[0] * m for _ in range(n)]  # Do not use [[0] * m] * n.
 57# indexed = list(enumerate(a))
 58# paired = list(zip(a, b))
 59# adjacent = list(pairwise(a))
 60# adjacent_compatible = list(zip(a, a[1:]))
 61
 62
 63# ============================================================
 64# 2. Pairs and intervals
 65# ============================================================
 66
 67def iter_pairs(n):
 68    """Yield all index pairs (i, j) with 0 <= i < j < n."""
 69    for i in range(n):
 70        for j in range(i + 1, n):
 71            yield i, j
 72
 73
 74def iter_intervals(n):
 75    """Yield all non-empty half-open intervals [left, right)."""
 76    # Note: In 0-indexed arrays, generating half-open intervals [left, right)
 77    # is mathematically identical to generating pairs (i, j) where i < j <= n.
 78    yield from iter_pairs(n + 1)
 79
 80
 81# ============================================================
 82# 3. Subsets and Cartesian products
 83# ============================================================
 84
 85def iter_subsets_mask(a):
 86    """Yield (mask, subset) for every subset of a."""
 87    n = len(a)
 88    for mask in range(1 << n):
 89        subset = tuple(a[i] for i in range(n) if mask >> i & 1)
 90        yield mask, subset
 91
 92
 93def iter_subsets(a):
 94    """Yield every subset as a tuple, grouped by subset size."""
 95    for size in range(len(a) + 1):
 96        yield from combinations(a, size)
 97
 98
 99def iter_binary_states(n):
100    """Yield all binary states of length n as tuples."""
101    yield from product([0, 1], repeat=n)
102
103
104def iter_k_states(k, n):
105    """Yield all k-ary states of length n as tuples."""
106    yield from product(range(k), repeat=n)
107
108
109# ============================================================
110# 4. Permutations and combinations
111# ============================================================
112
113def iter_multisets(a, k):
114    """Yield all multisets (combinations with replacement) of size k."""
115    yield from combinations_with_replacement(a, k)
116
117# for order in permutations(a):
118#     ...
119#
120# for chosen in combinations(a, k):
121#     ...
122#
123# for chosen in combinations_with_replacement(a, k):
124#     ...
125
126
127def unique_permutations(a):
128    """Yield distinct permutations even when a contains duplicates."""
129    items = sorted(a)
130    used = [False] * len(items)
131    path = []
132
133    def dfs():
134        if len(path) == len(items):
135            yield tuple(path)
136            return
137
138        for i, value in enumerate(items):
139            if used[i]:
140                continue
141            if i > 0 and items[i] == items[i - 1] and not used[i - 1]:
142                continue
143
144            used[i] = True
145            path.append(value)
146            yield from dfs()
147            path.pop()
148            used[i] = False
149
150    yield from dfs()
151
152
153# ============================================================
154# 5. DFS / backtracking
155# ============================================================
156
157# Example:
158# >>> dfs_assignments([["a","b"], [1,2]])
159# [('a', 1), ('a', 2), ('b', 1), ('b', 2)]
160def dfs_assignments(options):
161    """Return every sequence that chooses one value per position."""
162    answer = []
163    path = []
164
165    def dfs(position):
166        if position == len(options):
167            answer.append(tuple(path))
168            return
169
170        for choice in options[position]:
171            # Manually uncomment the next line when pruning is needed:
172            # if sum(path) + choice > SOME_LIMIT: continue
173            path.append(choice)
174            dfs(position + 1)
175            path.pop()
176
177    dfs(0)
178    return answer
179
180
181# ============================================================
182# 6. BFS shortest path in an implicit state graph
183# ============================================================
184
185# Example:
186# >>> def neighbors(x): return [y for y in (x-1, x+1) if 0 <= y <= 4]
187# >>> bfs_shortest(0, lambda x: x == 3, neighbors)
188# 3
189def bfs_shortest(start, is_goal, neighbors):
190    """Return the minimum number of edges to a goal, or None."""
191    queue = deque([start])
192    distance = {start: 0}
193
194    while queue:
195        state = queue.popleft()
196        current_distance = distance[state]
197
198        if is_goal(state):
199            return current_distance
200
201        for next_state in neighbors(state):
202            if next_state in distance:
203                continue
204            distance[next_state] = current_distance + 1
205            queue.append(next_state)
206
207    return None
208
209
210# ============================================================
211# 7. Memoized DFS example
212# ============================================================
213
214def subset_sum_exists(a, target):
215    """Return whether a subset of a sums to target."""
216    values = tuple(a)
217
218    @cache
219    def dfs(index, current_sum):
220        if index == len(values):
221            return current_sum == target
222
223        return (
224            dfs(index + 1, current_sum)
225            or dfs(index + 1, current_sum + values[index])
226        )
227
228    return dfs(0, 0)
229
230
231# ============================================================
232# 8. Prefix sums and small predicates
233# ============================================================
234
235def prefix_sums(a):
236    """Return [0, a[0], a[0]+a[1], ...]."""
237    return list(accumulate(a, initial=0))
238
239
240def range_sum(prefix, left, right):
241    """Return the sum on the half-open interval [left, right)."""
242    return prefix[right] - prefix[left]
243
244
245def is_strictly_increasing(a):
246    return all(x < y for x, y in pairwise(a))
247
248
249def is_square(n):
250    if n < 0:
251        return False
252    root = isqrt(n)
253    return root * root == n
254
255
256def first_true(candidates, predicate):
257    return next((x for x in candidates if predicate(x)), None)
258
259
260# ============================================================
261# 9. Containers and graphs
262# ============================================================
263
264def frequency(a):
265    return Counter(a)
266
267
268def group_by(items, key):
269    groups = defaultdict(list)
270    for item in items:
271        groups[key(item)].append(item)
272    return dict(groups)
273
274
275def build_undirected_graph(n, edges):
276    graph = [[] for _ in range(n)]
277    for u, v in edges:
278        assert 0 <= u < n and 0 <= v < n
279        graph[u].append(v)
280        graph[v].append(u)
281    return graph
282
283
284# State deduplication:
285# visited = set()
286# state = [1, 2, 3]
287# visited.add(tuple(state))
288
289
290# ============================================================
291# 10. Heap and binary search quick reference
292# ============================================================
293
294# heap = [5, 1, 4]
295# heapify(heap)
296# heappush(heap, 2)
297# smallest = heappop(heap)
298#
299# ordered = [1, 3, 3, 7]
300# first_three = bisect_left(ordered, 3)
301# after_three = bisect_right(ordered, 3)
302
303
304# ============================================================
305# 11. Replace this with the current problem
306# ============================================================
307
308def solve():
309    # Example:
310    # n, target = read_ints()
311    # a = read_ints()
312    # print("YES" if subset_sum_exists(a, target) else "NO")
313    pass
314
315
316# ============================================================
317# 12. Template self-test; delete after copying if not needed
318# ============================================================
319
320def _self_test():
321    from io import BytesIO
322
323    source = BytesIO(b"1 2 3\n4 5\n")
324    assert read_ints(source.readline) == [1, 2, 3]
325    assert read_all_ints(source.read) == [4, 5]
326
327    assert list(iter_pairs(0)) == []
328    assert list(iter_pairs(1)) == []
329    assert list(iter_pairs(3)) == [(0, 1), (0, 2), (1, 2)]
330    assert list(iter_intervals(0)) == []
331    assert list(iter_intervals(2)) == [(0, 1), (0, 2), (1, 2)]
332
333    subsets_mask = list(iter_subsets_mask([10, 20]))
334    assert subsets_mask == [
335        (0, ()),
336        (1, (10,)),
337        (2, (20,)),
338        (3, (10, 20)),
339    ]
340    assert list(iter_subsets([])) == [()]
341    assert list(iter_subsets([1, 2])) == [(), (1,), (2,), (1, 2)]
342
343    assert list(iter_binary_states(2)) == [
344        (0, 0),
345        (0, 1),
346        (1, 0),
347        (1, 1),
348    ]
349    
350    assert list(iter_k_states(3, 2)) == [
351        (0, 0), (0, 1), (0, 2),
352        (1, 0), (1, 1), (1, 2),
353        (2, 0), (2, 1), (2, 2)
354    ]
355
356    assert list(permutations([1, 2])) == [(1, 2), (2, 1)]
357    assert list(combinations([1, 2, 3], 2)) == [(1, 2), (1, 3), (2, 3)]
358    assert list(iter_multisets([1, 2], 2)) == [
359        (1, 1),
360        (1, 2),
361        (2, 2),
362    ]
363    assert list(unique_permutations([])) == [()]
364    assert list(unique_permutations([1, 1, 2])) == [
365        (1, 1, 2),
366        (1, 2, 1),
367        (2, 1, 1),
368    ]
369
370    assert dfs_assignments([]) == [()]
371    assert dfs_assignments([[0, 1], ["a", "b"]]) == [
372        (0, "a"),
373        (0, "b"),
374        (1, "a"),
375        (1, "b"),
376    ]
377
378    def line_neighbors(x):
379        return [y for y in (x - 1, x + 1) if 0 <= y <= 4]
380
381    assert bfs_shortest(0, lambda x: x == 0, line_neighbors) == 0
382    assert bfs_shortest(0, lambda x: x == 3, line_neighbors) == 3
383    assert bfs_shortest(0, lambda x: x == 1, lambda _x: ()) is None
384
385    assert subset_sum_exists([], 0)
386    assert subset_sum_exists([2, 3, 7], 5)
387    assert not subset_sum_exists([2, 4], 5)
388
389    prefix = prefix_sums([3, -2, 5, -1])
390    assert prefix == [0, 3, 1, 6, 5]
391    assert range_sum(prefix, 1, 3) == 3
392    assert is_strictly_increasing([])
393    assert is_strictly_increasing([1, 3, 8])
394    assert not is_strictly_increasing([1, 3, 3])
395    assert is_square(0)
396    assert is_square(10**20)
397    assert not is_square(15)
398    assert not is_square(-1)
399    assert first_true(range(10), lambda x: x > 5) == 6
400    assert first_true(range(3), lambda x: x > 5) is None
401
402    assert frequency([1, 1, 2]) == Counter({1: 2, 2: 1})
403    assert group_by([1, 2, 3, 4], lambda x: x % 2) == {
404        1: [1, 3],
405        0: [2, 4],
406    }
407    assert build_undirected_graph(3, [(0, 1), (1, 2)]) == [
408        [1],
409        [0, 2],
410        [1],
411    ]
412
413    heap = [5, 1, 4]
414    heapify(heap)
415    heappush(heap, 2)
416    assert [heappop(heap) for _ in range(4)] == [1, 2, 4, 5]
417
418    ordered = [1, 3, 3, 7]
419    assert bisect_left(ordered, 3) == 1
420    assert bisect_right(ordered, 3) == 3
421    assert gcd(18, 24) == 6
422    assert lcm(6, 8) == 24
423    assert INF > 10**100
424
425    print("brute_force_template: self-test passed")
426
427
428if __name__ == "__main__":
429    if "--self-test" in sys.argv:
430        _self_test()
431    else:
432        solve()

1def solve():
2    # n, target = read_ints()
3    # a = read_ints()
4    # print("YES" if subset_sum_exists(a, target) else "NO")
5    pass

14 9
22 7 11 15

 1def solve_text(text):
 2    lines = iter(text.strip().splitlines())
 3    n, target = map(int, next(lines).split())
 4    a = list(map(int, next(lines).split()))
 5    assert len(a) == n
 6
 7    answer = any(
 8        sum(a[i] for i in range(n) if mask >> i & 1) == target
 9        for mask in range(1 << n)
10    )
11    return "YES" if answer else "NO"
12
13
14assert solve_text("4 9\n2 7 11 15\n") == "YES"
15assert solve_text("3 20\n2 7 11\n") == "YES"
16assert solve_text("3 100\n2 7 11\n") == "NO"

1# n, target = read_ints()
2# a = read_ints()

 1def iter_subsets_mask(a):
 2    n = len(a)
 3    for mask in range(1 << n):
 4        subset = tuple(a[i] for i in range(n) if mask >> i & 1)
 5        yield mask, subset
 6
 7
 8subsets = list(iter_subsets_mask([10, 20]))
 9
10assert subsets == [
11    (0, ()),
12    (1, (10,)),
13    (2, (20,)),
14    (3, (10, 20)),
15]

1from itertools import combinations
2
3
4def iter_subsets(a):
5    for size in range(len(a) + 1):
6        yield from combinations(a, size)
7
8
9assert list(iter_subsets([1, 2])) == [(), (1,), (2,), (1, 2)]

 1from itertools import permutations
 2
 3
 4def minimum_adjacent_cost(a):
 5    return min(
 6        sum(abs(order[i] - order[i + 1]) for i in range(len(order) - 1))
 7        for order in permutations(a)
 8    )
 9
10
11assert minimum_adjacent_cost([1, 4, 6]) == 5

1bfs_shortest(start, is_goal, neighbors)

 1def make_integer_bfs(target):
 2    def is_goal(x):
 3        return x == target
 4
 5    def neighbors(x):
 6        for next_x in (x - 1, x + 1, x * 2):
 7            if 0 <= next_x <= 100:
 8                yield next_x
 9
10    return is_goal, neighbors
11
12
13is_goal, neighbors = make_integer_bfs(17)
14
15assert not is_goal(5)
16assert set(neighbors(5)) == {4, 6, 10}

1# answer = bfs_shortest(5, is_goal, neighbors)
2# print(answer)

 1def dfs_assignments(options):
 2    answer = []
 3    path = []
 4
 5    def dfs(position):
 6        if position == len(options):
 7            answer.append(tuple(path))
 8            return
 9
10        for choice in options[position]:
11            if choice in path:  # 示例剪枝：不允许重复选择
12                continue
13            path.append(choice)
14            dfs(position + 1)
15            path.pop()
16
17    dfs(0)
18    return answer
19
20
21assert dfs_assignments([[1, 2], [1, 2]]) == [(1, 2), (2, 1)]

 1from itertools import product
 2
 3assert list(product([0, 1], repeat=2)) == [
 4    (0, 0),
 5    (0, 1),
 6    (1, 0),
 7    (1, 1),
 8]
 9
10assert len(list(product(range(3), repeat=2))) == 3**2

1from itertools import accumulate
2
3a = [3, -2, 5, -1]
4prefix = list(accumulate(a, initial=0))
5
6left, right = 1, 3
7assert prefix[right] - prefix[left] == sum(a[left:right]) == 3

 1from collections import Counter, defaultdict
 2
 3a = [1, 2, 1, 3, 2]
 4count = Counter(a)
 5groups = defaultdict(list)
 6
 7for x in a:
 8    groups[x % 2].append(x)
 9
10assert count == Counter({1: 2, 2: 2, 3: 1})
11assert groups[0] == [2, 2]
12assert groups[1] == [1, 1, 3]

1python3 content/program_language/python/src/brute_force_template.py

1python3 content/program_language/python/src/brute_force_template.py --self-test

1python3 -m py_compile content/program_language/python/src/brute_force_template.py