Sosic and Gu n皇后算法实现：O(1) partial_collision函数优化问题

我正在实现用于n皇后问题的Sosic and Gu算法，其中initial_search()是初始化阶段。算法逻辑如下：

• 在int(3.08 * n)次迭代限制内，随机放置皇后到棋盘上
• 每次放置尝试时，调用partial_collision函数检查与已放置皇后的即时冲突
• 若检测到冲突，则重新随机选择位置，直到找到有效放置
• 这一步的目标是让皇后在棋盘上合理分布，最小化初始冲突

目前卡在partial_collision()的实现上：论文指出该函数需要达到**O(1)**时间复杂度，但我找不到利用对角数组仅检查左侧列的方法（已经实现了O(1)的total_collision()）。

以下是我的代码：

import random


def queen_search(queens):
    while True:
        nd, pd = initialize_diagonal_arrays(len(queens))
        k = initial_search(queens, nd, pd)
        if len(queens) > 200:
            final_search(queens, k, nd, pd)
            break
        else:
            final_search_reduced(queens, k, nd, pd)
            if board_collision(queens) == 0:
                break


def initial_search(queens, nd, pd):
    n = len(queens)
    queens[:] = list(range(n))
    update_diagonal_arrays(queens, nd, pd)
    j = 0

    # place queens without collisions
    for i in range(int(3.08 * n)):
        if j == n:
            break
        m = random.randint(j, n - 1)
        swap(queens, j, m, nd, pd)
        if partial_collision(queens, j) == 0:
            j += 1
        else:
            swap(queens, j, m, nd, pd)

    # place queens with possible collisions
    for i in range(j, n):
        m = random.randint(i, n - 1)
        swap(queens, i, m, nd, pd)

    # return the number of queens with possible collisions
    return n - j


def final_search(queens, k, nd, pd):
    n = len(queens)
    it = 0
    for i in range(n - k, n):
        if total_collisions(queens, i, nd, pd) > 0:
            while it < 7000:
                j = random.randint(0, n - 1)
                swap(queens, i, j, nd, pd)
                b = (total_collisions(queens, i, nd, pd) > 0) or (total_collisions(queens, j, nd, pd) > 0)
                if b:
                    swap(queens, i, j, nd, pd)
                    it += 1
                else:
                    break


def final_search_reduced(queens, k, nd, pd):
    n = len(queens)
    for i in range(n - k, n):
        if total_collisions(queens, i, nd, pd) > 0:
            for j in range(n):
                swap(queens, i, j, nd, pd)
                b = (total_collisions(queens, i, nd, pd) > 0) or (total_collisions(queens, j, nd, pd) > 0)
                if b:
                    swap(queens, i, j, nd, pd)
                else:
                    break


# auxiliary functions
def initialize_diagonal_arrays(n):
    neg_diagonal = [0] * (2 * n - 1)
    pos_diagonal = [0] * (2 * n - 1)
    return neg_diagonal, pos_diagonal


def update_diagonal_arrays(queens, neg_diagonal, pos_diagonal):
    n = len(queens)
    for i in range(len(queens)):
        neg_diagonal[i + queens[i]] += 1
        pos_diagonal[i - queens[i] + n - 1] += 1


def swap(queens, a, b, neg_diagonal, pos_diagonal):
    original_a, original_b = queens[a], queens[b]
    queens[a], queens[b] = queens[b], queens[a]

    neg_diagonal[a + original_a] -= 1
    neg_diagonal[b + original_b] -= 1
    pos_diagonal[a - original_a + len(queens) - 1] -= 1
    pos_diagonal[b - original_b + len(queens) - 1] -= 1

    neg_diagonal[a + queens[a]] += 1
    neg_diagonal[b + queens[b]] += 1
    pos_diagonal[a - queens[a] + len(queens) - 1] += 1
    pos_diagonal[b - queens[b] + len(queens) - 1] += 1


def partial_collision(queens, i):
    return sum(1 for j in range(i) if (i - j) == abs(queens[i] - queens[j]))


def total_collisions(queens, i, neg_diagonal, pos_diagonal):
    n = len(queens)
    return neg_diagonal[i + queens[i]] + pos_diagonal[i - queens[i] + n - 1] - 2  # Exclude self-collision


def board_collision(queens):
    return sum(partial_collision(queens, i) for i in range(len(queens)))

内容的提问来源于stack exchange，提问作者Tacce