Python itertools 与 functools 模块

import itertools

# count - 无限计数
for i in itertools.count(10, 2):  # 从10开始，步长2
    if i > 20:
        break
    print(i, end=" ")  # 10 12 14 16 18 20
print()

# cycle - 无限循环
for i, char in enumerate(itertools.cycle("AB")):
    if i >= 6:
        break
    print(char, end=" ")  # A B A B A B
print()

# repeat - 重复
print(list(itertools.repeat(5, 3)))  # [5, 5, 5]

import itertools

# chain - 连接多个迭代器
for item in itertools.chain([1, 2], [3, 4], [5, 6]):
    print(item, end=" ")  # 1 2 3 4 5 6
print()

# compress - 按选择器过滤
names = ['Alice', 'Bob', 'Charlie', 'David']
selector = [True, False, True, False]
print(list(itertools.compress(names, selector)))
# ['Alice', 'Charlie']

# dropwhile / takewhile
numbers = [1, 2, 3, 4, 5, 1, 2]

# 跳过小于4的元素
print(list(itertools.dropwhile(lambda x: x < 4, numbers)))
# [4, 5, 1, 2]

# 取小于4的元素
print(list(itertools.takewhile(lambda x: x < 4, numbers)))
# [1, 2, 3]

# filterfalse - 反向过滤
print(list(itertools.filterfalse(lambda x: x % 2 == 0, range(10))))
# [1, 3, 5, 7, 9]

# groupby - 分组（需要先排序）
data = [('A', 1), ('A', 2), ('B', 3), ('B', 4), ('A', 5)]
data.sort(key=lambda x: x[0])  # 必须按分组键排序

for key, group in itertools.groupby(data, key=lambda x: x[0]):
    print(f"{key}: {list(group)}")
# A: [(A, 1), (A, 2), (A, 5)]
# B: [(B, 3), (B, 4)]

import itertools

items = ['A', 'B', 'C']

# product - 笛卡尔积
print("product:", list(itertools.product(items, repeat=2)))
# [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), ...]

# permutations - 排列
print("permutations:", list(itertools.permutations(items, 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

# combinations - 组合
print("combinations:", list(itertools.combinations(items, 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'C')]

# combinations_with_replacement - 可重复组合
print("combinations_r:", list(itertools.combinations_with_replacement(items, 2)))
# [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

import itertools

# 生成扑克牌
suits = ['♠', '♥', '♣', '♦']
ranks = ['A', '2', '3', '4', '5', '6', '7', '8', '9', '10', 'J', 'Q', 'K']

deck = list(itertools.product(suits, ranks))
print(f"共 {len(deck)} 张牌")
print("前5张:", deck[:5])

# 批量处理
def batch(iterable, size):
    """将可迭代对象分批"""
    it = iter(iterable)
    while batch := list(itertools.islice(it, size)):
        yield batch

numbers = range(10)
for i, b in enumerate(batch(numbers, 3)):
    print(f"批次 {i}: {b}")
# 批次 0: [0, 1, 2]
# 批次 1: [3, 4, 5]
# 批次 2: [6, 7, 8]
# 批次 3: [9]

from functools import partial

# 固定部分参数
base_two = partial(int, base=2)
print(base_two("1010"))   # 10（二进制转十进制）
print(base_two("1111"))   # 15

# 固定更多参数
def power(base, exponent):
    return base ** exponent

square = partial(power, exponent=2)
cube = partial(power, exponent=3)

print(square(5))  # 25
print(cube(3))    # 27

# 实际应用：固定 print 参数
print_info = partial(print, "[INFO]", sep=" ")
print_info("程序启动")   # [INFO] 程序启动
print_info("处理完成")   # [INFO] 处理完成

from functools import lru_cache

# 无缓存
# def fibonacci(n):
#     if n < 2:
#         return n
#     return fibonacci(n - 1) + fibonacci(n - 2)

# 有缓存
@lru_cache(maxsize=128)
def fibonacci(n):
    if n < 2:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

print(fibonacci(100))  # 非常快！

# 查看缓存信息
print(fibonacci.cache_info())
# CacheInfo(hits=98, misses=101, maxsize=128, currsize=101)

# 清空缓存
fibonacci.cache_clear()

from functools import wraps

def my_decorator(func):
    @wraps(func)  # 保留原函数元信息
    def wrapper(*args, **kwargs):
        """包装函数"""
        print("执行前")
        result = func(*args, **kwargs)
        print("执行后")
        return result
    return wrapper

@my_decorator
def greet(name):
    """问候函数"""
    return f"你好, {name}"

# 使用 @wraps 后
print(greet.__name__)   # greet（不是 wrapper）
print(greet.__doc__)    # 问候函数（不是 包装函数）

from functools import cmp_to_key

# 自定义比较函数
def compare(x, y):
    """按绝对值比较"""
    if abs(x) < abs(y):
        return -1
    elif abs(x) > abs(y):
        return 1
    else:
        return 0

numbers = [3, -1, -4, 1, -5, 9, -2]
sorted_nums = sorted(numbers, key=cmp_to_key(compare))
print(sorted_nums)  # [-1, 1, -2, 3, -4, -5, 9]

from functools import reduce

# 求和
numbers = [1, 2, 3, 4, 5]
print(reduce(lambda x, y: x + y, numbers))  # 15

# 连乘
print(reduce(lambda x, y: x * y, numbers))  # 120

# 找最大值
print(reduce(lambda x, y: x if x > y else y, numbers))  # 5

# 字符串拼接
words = ["Hello", " ", "World"]
print(reduce(lambda x, y: x + y, words))  # Hello World

from functools import total_ordering

@total_ordering
class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age
    
    def __eq__(self, other):
        return self.age == other.age
    
    def __lt__(self, other):
        return self.age < other.age

# 只需实现 __eq__ 和 __lt__，其他比较自动获得
alice = Person("Alice", 25)
bob = Person("Bob", 30)

print(alice < bob)   # True
print(alice <= bob)  # True
print(alice > bob)   # False
print(alice >= bob)  # False
print(alice == bob)  # False
print(alice != bob)  # True