小C的w五元组问题解析 | 豆包MarsCode AI刷题

题目详解：小C的w五元组问题

问题理解

小C拿到了一个长度为 n 的数组 a_1, a_2, ..., a_n，她想知道有多少组 (i, j, k, h, l) 为 "w五元组"。w五元组定义如下：

• 满足条件的下标必须满足 1 ≤ i < j < k < h < l ≤ n；
• 数组元素需要满足：a_i = a_k = a_l，且 a_j = a_h，并且 a_i > a_j。

目标是求出所有满足条件的w五元组数量，并对结果取模 10^9 + 7。

数据结构选择

为了高效地解决这个问题，我们需要选择合适的数据结构来存储和查找数组元素的索引。具体来说，我们可以使用哈希表（HashMap）来存储每个元素的索引列表。这样可以在常数时间内查找某个元素的所有索引，从而加速查找过程。

算法步骤

1. 初始化哈希表：遍历数组，将每个元素的索引存储在哈希表中。
2. 遍历数组：对于每个元素 a_i，查找所有可能的 (i, j, k, h, l) 组合。
3. 嵌套循环：使用嵌套循环来检查每个组合是否满足条件：
   • 确保 1 ≤ i < j < k < h < l ≤ n。
   • 确保 a_i = a_k = a_l，且 a_j = a_h，并且 a_i > a_j。
4. 计数并取模：在每次找到满足条件的五元组时，计数加1，并对结果取模 10^9 + 7。

代码详解

import java.util.*;

public class Main {
    public static int solution(int n, List<Integer> a) {
        final int MOD = 1_000_000_007;
        Map<Integer, List<Integer>> indexMap = new HashMap<>();

        // 将每个元素的索引存储在哈希表中
        for (int i = 0; i < n; i++) {
            indexMap.computeIfAbsent(a.get(i), k -> new ArrayList<>()).add(i);
        }

        int count = 0;

        // 遍历每个元素，查找满足条件的五元组
        for (int i = 0; i < n; i++) {
            int ai = a.get(i);
            List<Integer> indices = indexMap.get(ai);

            for (int j = i + 1; j < n; j++) {
                int aj = a.get(j);
                if (ai <= aj) continue; // 不满足 a_i > a_j 的条件

                for (int k : indices) {
                    if (k <= j) continue; // 不满足 i < j < k 的条件

                    for (int h = k + 1; h < n; h++) {
                        int ah = a.get(h);
                        if (ah != aj) continue; // 不满足 a_j = a_h 的条件

                        for (int l : indices) {
                            if (l <= h) continue; // 不满足 k < h < l 的条件
                            count = (count + 1) % MOD;
                        }
                    }
                }
            }
        }

        return count;
    }

    public static void main(String[] args) {
        System.out.println(solution(7, Arrays.asList(3, 1, 3, 1, 3, 1, 3)) == 6);
        System.out.println(solution(6, Arrays.asList(2, 1, 2, 1, 2, 1)) == 1);
        System.out.println(solution(5, Arrays.asList(5, 3, 5, 3, 5)) == 1);
    }
}

代码解析

1. 哈希表 indexMap：

   • indexMap 是一个 HashMap，键为数组元素，值为该元素在数组中的索引列表。
   • computeIfAbsent 方法用于在哈希表中不存在某个键时，初始化一个空列表，并将当前索引添加到该列表中。

2. 遍历数组：

   • 外层循环遍历数组中的每个元素 a_i。
   • 内层循环遍历 a_i 之后的元素 a_j，确保 a_i > a_j。

3. 嵌套循环：

   • 对于每个 a_i，查找所有可能的 (i, j, k, h, l) 组合。
   • 使用 indices 列表来查找 a_i 的所有索引，确保 i < j < k。
   • 对于每个 k，查找 a_j 的所有索引，确保 k < h < l。

4. 计数并取模：

   • 在每次找到满足条件的五元组时，计数加1，并对结果取模 10^9 + 7，防止溢出。

测试样例解析

1. 样例1：

   • 输入：n = 7 , a = [3, 1, 3, 1, 3, 1, 3]
   • 输出：6
   • 解析：满足条件的五元组有 (0, 1, 2, 3, 4)、(0, 1, 2, 3, 6)、(0, 1, 4, 3, 6)、(0, 1, 4, 5, 6)、(2, 3, 4, 5, 6)、(0, 1, 2, 5, 6)。

2. 样例2：

   • 输入：n = 6 , a = [2, 1, 2, 1, 2, 1]
   • 输出：1
   • 解析：满足条件的五元组只有 (0, 1, 2, 3, 4)。

3. 样例3：

   • 输入：n = 5 , a = [5, 3, 5, 3, 5]
   • 输出：1
   • 解析：满足条件的五元组只有 (0, 1, 2, 3, 4)。

总结

通过使用哈希表来存储每个元素的索引列表，我们可以在常数时间内查找某个元素的所有索引，从而加速查找过程。嵌套循环的使用虽然增加了时间复杂度，但在合理的数据结构选择下，仍然可以高效地解决问题。最终，我们在每次找到满足条件的五元组时，计数加1，并对结果取模 10^9 + 7，防止溢出。