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