排列组合查询算法优化：如何降低O(mn+qm)的时间复杂度？

问题描述 #

给定数组 $a_1, \dots, a_n$，每个元素是长度为 $m$ 的排列。两个排列 $x$ 和 $y$ 的复合操作定义为 $z = x * y$，其中 $z(i) = y(x(i))$。给定 $q$ 个查询，每个查询包含两个整数 $1 \leq l \leq r \leq n$，令 $b = ((\dots((a_l*a_{l+1})a_{l+2})\dots)a_r)$，查询答案为 $\text{sum} = 1b(1)+\dots+mb(m)$，需要实现快速响应查询的程序。

约束条件 #

• $1 \leq n,m \leq 10^5$，且 $1 \leq nm \leq 210^5$
• $1 \leq q \leq 2*10^5$
• 所有测试用例的 $n \cdot m$ 和 $q$ 之和不超过 $2*10^5$
• Python 时间限制 6 秒，C++ 为 2 秒，内存限制 512MB

当前解法与问题 #

我采用前缀排列数组与前缀排列逆数组的思路实现，但代码时间复杂度为 $O(mn+qm)$，运行速度过慢，请问如何优化？

我的Python代码 #

t = int(input())
for i in range (t):
    n,m = list(map(int,input().split()))
    prefix_mass = []
    inv_mass = []
    s = list(range(0, m))
    s_ = list(range(0, m))
    for j in range (n):
        e = list(map(int,input().split()))
        new_s = [0] * m
        new_s_ = [0] * m
        a_ = [0] * m
        for k in range (m):
            new_s[k] = e[s[k]] - 1
            new_s_[e[k] - 1] = s_[k]
        s = new_s
        s_ = new_s_
        prefix_mass.append(s)
        inv_mass.append(s_)
    q = int(input())
    for j in range(q):
        l, r = list(map(int,input().split()))
        l -= 1
        r -= 1
        if l == 0:
            ans = 0
            s = prefix_mass[r]
            for k in range(m):
                ans += (s[k] + 1) * (k + 1)
            if (i == t - 1 and j == q - 1):
                print(ans, end="")
            else: print(ans)
        else:
            ans = 0
            s = prefix_mass[r]
            s_ = inv_mass[l - 1]
            for k in range(m):
                ans += (k + 1) * (s[s_[k]] + 1)
            if (i == t - 1 and j == q - 1):
                print(ans, end="")
            else: print(ans)

输入输出示例 #

输入 #

1
4 3
3 2 1
1 3 2
1 2 3
2 3 1
5
1 1
1 4
3 4
3 3
1 3

输出 #

10
11
11
14
11

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