python 算到16皇后又沒有可能不爆tle

python 算到16皇后又沒有可能不爆tle

過了

import sys

sys.setrecursionlimit(10000)

def count_n_queens(n: int) -> int:
    if n == 0:
        return 1
    if n < 0:
        return 0
    all_ones = (1 << n) - 1

    def backtrack(cols: int, diag_l: int, diag_r: int) -> int:
        if cols == all_ones:
            return 1
        cnt = 0
        avail = all_ones & ~(cols | diag_l | diag_r)
        while avail:
            bit = avail & -avail
            avail -= bit
            cnt += backtrack(cols | bit, ((diag_l | bit) << 1) & all_ones, (diag_r | bit) >> 1)
        return cnt

    total = 0
    half = n // 2
    for c in range(half):
        b = 1 << c
        total += backtrack(b, (b << 1) & all_ones, b >> 1)
    total *= 2
    if n % 2 == 1:
        center = 1 << half
        total += backtrack(center, (center << 1) & all_ones, center >> 1)
    return total

def main():
    data = sys.stdin.buffer.read().split()
    if not data:
        return
    ns = []
    for token in data:
        try:
            ns.append(int(token))
        except:
            ns.append(0)
    result_cache = {}
    uniq = set(ns)
    for n in uniq:
        result_cache[n] = count_n_queens(n)
    out = "\n".join(str(result_cache[n]) for n in ns)
    sys.stdout.write(out)

if __name__ == "__main__":
    main()