HJ98 喜欢切数组的红（dp

描述

小红有一个长度为 nn 的数组 {a1,a2,...,an}{a1,a2,...,an} ，她打算将数组切两刀变成三个非空子数组，使得每一个子数组中至少存在一个正数，且每个子数组的和都相等。

看起来不是很难，所以小红想让你求解，一共有多少种不同的切分方案。

输入描述：

第一行输入两个整数 n(3≦n≦2×105) 代表数组中的元素数量。

第二行输入 nn 个整数 a1,a2,...,an(−109≦ai≦109)代表数组元素。

输出描述：

在一行上输出一个整数，代表切分方案数。

示例1

输入：

3

3 3 3

输出：1
示例2

输入：
复制代码
6
1 1 4 5 1 4
输出：

0
示例3

输入：
复制代码
10
0 3 4 2 3 2 1 -1 3 4
输出：
复制代码
2

代码思路

时间复杂度❌ O(n³) 没过

java 复制代码

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
       int n= in.nextInt();
       int[] nums=new int[n];
       for(int i=0;i<n;i++){
           nums[i]=in.nextInt();
       }
        System.out.println(fun(nums));

    }
    public static int fun(int[] nums){
        if(nums.length<3){
            return 0;
        }
        int sum= Arrays.stream(nums).sum();
        if(sum%3!=0){
            return 0;
        }
        int[] dp =new int[nums.length];
        dp[0]=nums[0];
        List<Integer> l1=new ArrayList<>();
        if(dp[0]==sum/3){
            l1.add(0);
        }
        for(int i=1;i<nums.length;i++){
            dp[i]=dp[i-1]+nums[i];
            if(dp[i]==sum/3){
                l1.add(i);
            }
        }
        if(l1.isEmpty()){
            return 0;
        }
        int res=0;
        for(Integer x:l1){
            if(hasPositive(nums,0,x)){
                for(int i=x+1;i<nums.length;i++){
                    if(dp[i]==sum/3*2&&hasPositive(nums,x+1,i)&&hasPositive(nums,i+1,nums.length-1)){
                        res++;
                    }
                }
            }

        }
        return res;
    }

    // 检查数组的指定区间[start, end]是否至少有一个正数
    public static boolean hasPositive(int[] nums, int start, int end) {
        for (int i = start; i <= end; i++) {
            if (nums[i] > 0) {
                return true;
            }
        }
        return false;
    }
}

可通行代码

java 复制代码

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int[] nums = new int[n];
        for (int i = 0; i < n; i++) {
            nums[i] = in.nextInt();
        }
        System.out.println(fun(nums));
    }

    public static int fun(int[] nums) {
        int n = nums.length;
        if (n < 3) return 0;

        int sum = Arrays.stream(nums).sum();
        if (sum % 3 != 0) return 0;

        int part = sum / 3;
        int part2 = 2 * part;

        // ---------- 构建前缀和和前缀正数统计 ----------
        int[] prefix = new int[n];
        int[] pos = new int[n];  // pos[i]：0~i 之间正数个数

        prefix[0] = nums[0];
        pos[0] = nums[0] > 0 ? 1 : 0;

        for (int i = 1; i < n; i++) {
            prefix[i] = prefix[i - 1] + nums[i];
            pos[i] = pos[i - 1] + (nums[i] > 0 ? 1 : 0);
        }

        // ---------- 找到满足 prefix[i] = part 的 i（第一刀） ----------
        List<Integer> list1 = new ArrayList<>();
        for (int i = 0; i < n - 2; i++) {  // i 最晚只能切到 n-3
            if (prefix[i] == part && hasPositive(pos, 0, i)) {
                list1.add(i);
            }
        }

        if (list1.isEmpty()) return 0;

        int res = 0;

        // ---------- 第二刀 ----------
        for (int i : list1) {
            for (int j = i + 1; j < n - 1; j++) {  // j 最晚只能切到 n-2
                if (prefix[j] == part2 &&
                    hasPositive(pos, i + 1, j) &&
                    hasPositive(pos, j + 1, n - 1)) {
                    res++;
                }
            }
        }

        return res;
    }

    // O(1) 判断区间是否存在正数
    public static boolean hasPositive(int[] pos, int l, int r) {
        if (l > r) return false;
        int left = l == 0 ? 0 : pos[l - 1];
        return pos[r] - left > 0;
    }
}