3618. 根据质数下标分割数组

题目描述

给你一个整数数组 nums。

根据以下规则将 nums 分割成两个数组 A 和 B：

nums 中位于质数下标的元素必须放入数组 A。
所有其他元素必须放入数组 B。

返回两个数组和的绝对差值：|sum(A) - sum(B)|。

质数是一个大于 1 的自然数，它只有两个因子，1和它本身。

注意：空数组的和为 0。

示例 1:

输入: nums = [2,3,4]

输出: 1

解释:

数组中唯一的质数下标是 2，所以 nums[2] = 4 被放入数组 A。
其余元素 nums[0] = 2 和 nums[1] = 3 被放入数组 B。
sum(A) = 4，sum(B) = 2 + 3 = 5。
绝对差值是 |4 - 5| = 1。

示例 2:

输入: nums = [-1,5,7,0]

输出: 3

解释:

数组中的质数下标是 2 和 3，所以 nums[2] = 7 和 nums[3] = 0 被放入数组 A。
其余元素 nums[0] = -1 和 nums[1] = 5 被放入数组 B。
sum(A) = 7 + 0 = 7，sum(B) = -1 + 5 = 4。
绝对差值是 |7 - 4| = 3。

提示:

1 <= nums.length <= 10⁵
-10⁹ <= nums[i] <= 10⁹

解法

方法一：埃氏筛 + 模拟

我们可以用埃氏筛法预处理出 $[0, 10^5]$ 范围内的所有质数。然后遍历数组 $ \textit{nums}$，对于 $\textit{nums}[i]$，如果 $i$ 是质数，则将 $\textit{nums}[i]$ 加到答案中，否则将 $-\textit{nums}[i]$ 加到答案中。最后返回答案的绝对值。

忽略预处理的时间和空间，时间复杂度 $O(n)$，其中 $n$ 是数组 $\textit{nums}$ 的长度，空间复杂度 $O(1)$。

Python3JavaC++GoTypeScript

m = 10**5 + 10
primes = [True] * m
primes[0] = primes[1] = False
for i in range(2, m):
    if primes[i]:
        for j in range(i + i, m, i):
            primes[j] = False


class Solution:
    def splitArray(self, nums: List[int]) -> int:
        return abs(sum(x if primes[i] else -x for i, x in enumerate(nums)))

class Solution {
    private static final int M = 100000 + 10;
    private static boolean[] primes = new boolean[M];

    static {
        for (int i = 0; i < M; i++) {
            primes[i] = true;
        }
        primes[0] = primes[1] = false;

        for (int i = 2; i < M; i++) {
            if (primes[i]) {
                for (int j = i + i; j < M; j += i) {
                    primes[j] = false;
                }
            }
        }
    }

    public long splitArray(int[] nums) {
        long ans = 0;
        for (int i = 0; i < nums.length; ++i) {
            ans += primes[i] ? nums[i] : -nums[i];
        }
        return Math.abs(ans);
    }
}

const int M = 1e5 + 10;
bool primes[M];
auto init = [] {
    memset(primes, true, sizeof(primes));
    primes[0] = primes[1] = false;
    for (int i = 2; i < M; ++i) {
        if (primes[i]) {
            for (int j = i + i; j < M; j += i) {
                primes[j] = false;
            }
        }
    }
    return 0;
}();

class Solution {
public:
    long long splitArray(vector<int>& nums) {
        long long ans = 0;
        for (int i = 0; i < nums.size(); ++i) {
            ans += primes[i] ? nums[i] : -nums[i];
        }
        return abs(ans);
    }
};

const M = 100000 + 10

var primes [M]bool

func init() {
    for i := 0; i < M; i++ {
        primes[i] = true
    }
    primes[0], primes[1] = false, false

    for i := 2; i < M; i++ {
        if primes[i] {
            for j := i + i; j < M; j += i {
                primes[j] = false
            }
        }
    }
}

func splitArray(nums []int) (ans int64) {
    for i, num := range nums {
        if primes[i] {
            ans += int64(num)
        } else {
            ans -= int64(num)
        }
    }
    return max(ans, -ans)
}

const M = 100000 + 10;
const primes: boolean[] = Array(M).fill(true);

const init = (() => {
    primes[0] = primes[1] = false;

    for (let i = 2; i < M; i++) {
        if (primes[i]) {
            for (let j = i + i; j < M; j += i) {
                primes[j] = false;
            }
        }
    }
})();

function splitArray(nums: number[]): number {
    let ans = 0;
    for (let i = 0; i < nums.length; i++) {
        ans += primes[i] ? nums[i] : -nums[i];
    }
    return Math.abs(ans);
}

3618. 根据质数下标分割数组

题目描述

解法

方法一：埃氏筛 + 模拟

评论