邻位交换的最小次数

运行时间: 116 ms

内存: 17.4 MB

from sortedcontainers import SortedList


def kthNextPermutation(
    nums: List[int], k: int, inplace=False, prev=False
) -> Optional[List[int]]:
    """下k个字典序的排列(可以存在重复元素)

    Args:
        nums: 原排列数组.
        k (int): 后续第k个(`本身算第0个`)字典序的排列.
        inplace (bool, optional): 是否原地修改. 默认为False.
        prev (bool, optional): 使用next还是prev. 默认使用next.

    Returns:
        Optional[List[int]]: 下k个字典序的排列,如果不存在,返回None.
    """
    if not nums:
        return
    counter = Counter([nums[-1]])
    sl = (
        SortedList([nums[-1]])
        if not prev
        else SortedList([nums[-1]], key=lambda x: -x)
    )

    fac = 1
    facPtr = 1
    curPerm = 0
    overlap = 1  # 重复元素的个数的乘积
    allPerm = 1  # 后缀里的所有排列个数
    for right in range(len(nums) - 2, -1, -1):
        if curPerm + k < allPerm:
            break
        num = nums[right]
        counter[num] += 1
        overlap *= counter[num]

        smaller = 0
        pos = sl.bisect_left(num)
        if pos == len(sl) or sl[pos] != num:  # set去重
            sl.add(num)
        for i, pre in enumerate(sl):
            if i >= pos:
                break
            # for pre in sl.islice(0, pos):
            smaller += (fac * counter[pre]) // overlap

        facPtr += 1
        fac *= facPtr
        curPerm += smaller
        allPerm = fac // overlap

    if curPerm + k >= allPerm:
        return

    res = []
    fac //= facPtr
    permCount = 0
    target = curPerm + k
    while permCount != target:
        for i, pre in enumerate(sl):
            curPerm = (fac * counter[pre]) // overlap  # 以当前元素开头的排列个数
            cand = permCount + curPerm
            if cand <= target:
                permCount = cand
                continue
            facPtr -= 1
            fac //= facPtr
            overlap //= counter[pre]
            res.append(pre)
            counter[pre] -= 1
            if not counter[pre]:
                sl.pop(i)
            break

    for pre in sl:
        res += [pre] * counter[pre]
    if inplace:
        nums[-len(res) :] = res
        return nums
    return nums[: len(nums) - len(res)] + res


class BIT1:
    """单点修改,区间和查询"""

    __slots__ = "size", "bit", "tree"

    def __init__(self, n: int):
        self.size = n
        self.bit = n.bit_length()
        self.tree = [0] * (n + 1)

    def add(self, index: int, delta: int) -> None:
        # index 必须大于0
        while index <= self.size:
            self.tree[index] += delta
            index += index & -index

    def _query(self, index: int) -> int:
        res = 0
        while index > 0:
            res += self.tree[index]
            index -= index & -index
        return res

    def query(self, left: int, right: int) -> int:
        return self._query(right) - self._query(left - 1)


class Solution:
    def getMinSwaps(self, num: str, k: int) -> int:
        n = len(num)
        arr = [*map(int, num)]
        brr = arr.copy()
        brr = kthNextPermutation(brr, k)
        idx = defaultdict(list)
        for i, b in enumerate(brr, 2):
            idx[b].append(i)
        cnt = defaultdict(int)
        al = []
        for a in arr:
            al.append(idx[a][cnt[a]])
            cnt[a] += 1
        bit = BIT1(n + 10)
        res = 0
        for i in al:
            res += bit.query(i + 1, n + 5)
            bit.add(i, 1)
        return res