生成满足源元素最小间隔要求的随机排列的优化方案咨询

给定包含n个唯一元素的序列a，需生成a的一个随机排列序列b，要求将b追加到a后，所有重复元素之间的最小距离不小于指定值d。

例如，当a=[1,2,3]且d=2时，6种可能的排列中仅前4种符合要求：

a         b
[1, 2, 3] (1, 2, 3) mindist = 3
[1, 2, 3] (1, 3, 2) mindist = 2
[1, 2, 3] (2, 1, 3) mindist = 2
[1, 2, 3] (2, 3, 1) mindist = 2
[1, 2, 3] (3, 1, 2) mindist = 1
[1, 2, 3] (3, 2, 1) mindist = 1

后两种的最小距离1 < d，不符合要求。

原实现代码如下：

import random
n = 10
alist = list(range(n))

blist = alist[:]

d = n//2

avail_indices = list(range(n))
for a_ind, a_val in enumerate(reversed(alist)):
  min_ind = max(d - a_ind - 1, 0)
  new_ind = random.choice(avail_indices[min_ind:])
  avail_indices.remove(new_ind)
  blist[new_ind] = a_val
print(alist, blist)

该实现的时间复杂度为O(n²)，核心瓶颈在于avail_indices.remove(new_ind)操作：列表的remove方法需要遍历查找元素，每次耗时O(n)，循环n次后总复杂度达到O(n²)，n增大时耗时会显著上升。

方案一：使用SortedList（第三方库，最优复杂度） #

借助sortedcontainers.SortedList可以实现O(logn)时间的插入、删除和二分查找操作，彻底解决原代码的性能瓶颈。

代码示例：

import random
from sortedcontainers import SortedList

n = 10
alist = list(range(n))
blist = [0] * n
d = n // 2

avail_indices = SortedList(range(n))
for a_ind, a_val in enumerate(reversed(alist)):
    min_pos = max(d - a_ind - 1, 0)
    # 找到可选范围的起始位置在有序列表中的索引
    start_idx = avail_indices.bisect_left(min_pos)
    # 在可选范围内随机选择一个元素的索引
    selected_idx = random.randint(start_idx, len(avail_indices) - 1)
    new_ind = avail_indices.pop(selected_idx)
    blist[new_ind] = a_val

print(alist, blist)

复杂度分析 #

每次bisect_left和pop操作的时间复杂度为O(logn)，循环n次后总时间复杂度为O(nlogn)，相比原O(n²)实现有数量级的性能提升，n越大优化效果越明显。

方案二：纯Python实现（无第三方库，近似高效） #

如果无法使用第三方库，可借助bisect模块优化查找逻辑，同时用索引pop替代元素remove，虽然理论复杂度仍为O(n²)，但实际运行效率远高于原代码。

代码示例：

import random
import bisect

n = 10
alist = list(range(n))
blist = [0] * n
d = n // 2

avail_indices = list(range(n))
for a_ind, a_val in enumerate(reversed(alist)):
    min_ind = max(d - a_ind - 1, 0)
    # 找到可用索引中第一个大于等于min_ind的位置
    start_pos = bisect.bisect_left(avail_indices, min_ind)
    # 在可选范围内随机选择一个索引位置
    selected_pos = random.randint(start_pos, len(avail_indices)-1)
    # 通过索引直接删除，比remove更高效
    new_ind = avail_indices.pop(selected_pos)
    blist[new_ind] = a_val

print(alist, blist)

方案三：线段树实现（纯Python，严格O(nlogn)） #

对于超大规模的n，可通过实现线段树来维护可用位置，支持O(logn)时间的第k个可用位置查询和更新操作，达到严格的O(nlogn)时间复杂度。

代码示例：

import random

class SegmentTree:
    def __init__(self, size):
        self.n = 1
        while self.n < size:
            self.n <<= 1
        self.tree = [0] * (2 * self.n)
        # 初始化叶子节点：标记所有位置为可用
        for i in range(size):
            self.tree[self.n + i] = 1
        # 构建线段树
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = self.tree[2*i] + self.tree[2*i+1]
    
    def query_kth(self, k):
        # 查询第k个可用位置（k从0开始计数）
        node = 1
        while node < self.n:
            left_count = self.tree[2*node]
            if k < left_count:
                node = 2*node
            else:
                k -= left_count
                node = 2*node + 1
        return node - self.n
    
    def update(self, pos):
        # 标记指定位置为已用
        node = self.n + pos
        self.tree[node] = 0
        node >>= 1
        while node >= 1:
            prev_val = self.tree[node]
            self.tree[node] = self.tree[2*node] + self.tree[2*node+1]
            if self.tree[node] == prev_val:
                break
            node >>= 1

def get_prefix_sum(st, r):
    # 计算[0, r]范围内的可用位置数量
    res = 0
    l_node = st.n + 0
    r_node = st.n + r + 1
    while l_node < r_node:
        if l_node % 2 == 1:
            res += st.tree[l_node]
            l_node += 1
        if r_node % 2 == 1:
            r_node -= 1
            res += st.tree[r_node]
        l_node >>= 1
        r_node >>= 1
    return res

n = 10
alist = list(range(n))
blist = [0] * n
d = n // 2

st = SegmentTree(n)
available_count = n

for a_ind, a_val in enumerate(reversed(alist)):
    min_ind = max(d - a_ind - 1, 0)
    # 计算小于min_ind的可用位置数量
    count_less_min = get_prefix_sum(st, min_ind-1) if min_ind > 0 else 0
    # 在可选范围内随机选择第k个可用位置
    selected_k = random.randint(count_less_min, available_count - 1)
    new_ind = st.query_kth(selected_k)
    blist[new_ind] = a_val
    st.update(new_ind)
    available_count -= 1

print(alist, blist)

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