Chương trình tìm tổng các số tại vị trí hoán vị hợp lệ trong Python

Cho số n và một "hoán vị đặc biệt" (mảng input_arr) độ dài n chứa các số từ 1 đến n nhưng một số vị trí bị thay bằng 0 (giá trị quên). Cần tìm tất cả các hoán vị hợp lệ khớp với mẫu này (các vị trí khác 0 phải giữ nguyên), lấy số thứ tự (1-based) của các hoán vị đó trong danh sách tất cả hoán vị sắp xếp theo thứ tự từ điển, và trả về tổng các số thứ tự đó modulo 10⁹+7.

Ví dụ: input_arr = [0, 2, 0], n = 3 → kết quả 7. Hai hoán vị hợp lệ: [1,2,3] (số 2) và [3,2,1] (số 5), tổng = 7.

Thuật toán giải quyết

1. Nếu không có 0 trong mảng: tính trực tiếp thứ tự từ điển của hoán vị
2. Nếu có cnt số 0:
   • Xác định các số còn thiếu (miss)
   • Tính tổng chỉ số các số thiếu (tmp)
   • Duyệt mảng, dùng Fenwick tree / binary search (bisect) để đếm số hoán vị nhỏ hơn tại mỗi vị trí
   • Kết hợp toán học tổ hợp với modulo để tính tổng thứ tự

Ví dụ minh họa

import bisect

def solve(input_arr, n):
    modulo = 10**9 + 7
    i2 = pow(2, modulo-2, modulo)
    fact = [1]
    for x in range(1, n+1):
        fact.append(fact[-1] * x % modulo)
    cnt = input_arr.count(0)
    if not cnt:
        res = 0
        seen_list = []
        for i, x in enumerate(input_arr, 1):
            tmp_val = bisect.bisect(seen_list, x)
            res += fact[n-i] * (x - 1 - tmp_val)
            res %= modulo
            seen_list.insert(tmp_val, x)
        return res + 1
    else:
        ik = pow(cnt, modulo-2, modulo)
        miss = [True] * n
        for x in input_arr:
            if x != 0:
                miss[x-1] = False
        miss_srtd = []
        tmp = 0
        for i, x in enumerate(miss, 1):
            if x:
                miss_srtd.append(i)
                tmp += i
        pre = [0]
        for x in miss:
            pre.append(pre[-1] + x)
        cnt_cu = 0
        s = tmp % modulo * ik % modulo
        srtdw = []
        res = z = 0
        for i, x in enumerate(input_arr, 1):
            if x:
                l = tmp_val = bisect.bisect(srtdw, x)
                l += z * bisect.bisect(miss_srtd, x) % modulo * ik % modulo
                p = x - 1 - l
                p *= fact[cnt]
                p %= modulo
                srtdw.insert(tmp_val, x)
                cnt_cu += cnt - pre[x]
            else:
                l = cnt_cu
                l *= ik
                l += z * i2 % modulo
                p = s - 1 - l
                p *= fact[cnt]
                p %= modulo
                z += 1
                res += p * fact[n-i] % modulo
                res %= modulo
        return (res + fact[cnt]) % modulo

print(solve([0, 2, 0], 3))

Đầu vào

[0, 2, 0], 3

Đầu ra

7