4051. Count Subarrays with Distant Sums

Description

You are given an integer array nums and two integers goal and k.

A subarray nums[i..j] is considered distant if the absolute difference between its sum and goal is at least k.

Return the number of distant subarrays.

Example 1:

Input: nums = [1,2,1], goal = 4, k = 1

Output: 5

Explanation:

The distant subarrays for k = 1 are:

`i`	`j`	`nums[i..j]`	Sum	`abs(sum - goal)`
0	0	`[1]`	1	3
1	1	`[2]`	2	2
2	2	`[1]`	1	3
0	1	`[1, 2]`	3	1
1	2	`[2, 1]`	3	1

Thus, the answer is 5.

Example 2:

Input: nums = [2,-1,3], goal = 2, k = 2

Output: 2

Explanation:

The distant subarrays for k = 2 are:

`i`	`j`	`nums[i..j]`	Sum	`abs(sum - goal)`
1	1	`[-1]`	-1	3
0	2	`[2, -1, 3]`	4	2

Thus, the answer is 2.

Example 3:

Input: nums = [-3,1,2], goal = 0, k = 3

Output: 2

Explanation:

The distant subarrays for k = 3 are:

`i`	`j`	`nums[i..j]`	Sum	`abs(sum - goal)`
0	0	`[-3]`	-3	3
1	2	`[1, 2]`	3	3

Thus, the answer is 2.

Constraints:

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

Solutions

Solution 1: Prefix Sum + Binary Indexed Tree

Thinking

There are quadratically many subarrays, so $n = 10^5$ rules out enumeration. The complement of $|sum - \textit{goal}| \ge k$ is $|sum - \textit{goal}| < k$; subtracting that count from the total is cleaner.

Prefix sums turn a subarray sum into a difference of two points. For each right endpoint we need how many earlier prefix sums fall inside a numeric interval.

After sorting the prefix sums we binary-search the Fenwick indices, query the interval, then insert the current value.

Let $s$ be the prefix-sum array of $\textit{nums}$ ($s[0] = 0$). The sum of the subarray $\textit{nums}[L..R-1]$ is $s[R] - s[L]$, and it is distant if and only if $|s[R] - s[L] - \textit{goal}| \ge k$.

There are $\frac{n(n+1)}{2}$ nonempty subarrays. We count those that fail the condition, i.e. $|s[R] - s[L] - \textit{goal}| < k$, and subtract that count from the total.

The inequality is equivalent to

\[s[R] - \textit{goal} - k < s[L] < s[R] - \textit{goal} + k\]

that is, $s[L]$ lies in the closed interval $[s[R] - \textit{goal} - k + 1,\, s[R] - \textit{goal} + k - 1]$.

Enumerate the prefix sums $v = s[R]$ from left to right. Among the prefix sums already inserted, query how many fall in $[a, b]$ and subtract that from the answer, then insert $v$. For discretization we sort $s$ and locate Binary Indexed Tree indices by binary search.

The time complexity is $O(n \times \log n)$ and the space complexity is $O(n)$, where $n$ is the length of $\textit{nums}$.

class BinaryIndexedTree {
    private final int n;
    private final long[] c;

    BinaryIndexedTree(int n) {
        this.n = n;
        this.c = new long[n + 1];
    }

    void update(int x, long delta) {
        while (x <= n) {
            c[x] += delta;
            x += x & -x;
        }
    }

    long query(int x) {
        long s = 0;
        while (x > 0) {
            s += c[x];
            x -= x & -x;
        }
        return s;
    }
}

class Solution {
    public long distantSubarrays(int[] nums, int goal, int k) {
        int n = nums.length;
        long[] s = new long[n + 1];

        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + nums[i];
        }

        long[] st = s.clone();
        Arrays.sort(st);

        long ans = (long) n * (n + 1) / 2;
        BinaryIndexedTree bit = new BinaryIndexedTree(st.length + 1);

        for (long v : s) {
            long a = v - goal - k + 1L;
            long b = v - goal + k - 1L;

            int l = lowerBound(st, a) + 1;
            int r = lowerBound(st, b + 1);

            if (l <= r) {
                ans -= bit.query(r) - bit.query(l - 1);
            }

            bit.update(lowerBound(st, v) + 1, 1);
        }

        return ans;
    }

    private int lowerBound(long[] nums, long target) {
        int l = 0;
        int r = nums.length;

        while (l < r) {
            int m = (l + r) >>> 1;
            if (nums[m] < target) {
                l = m + 1;
            } else {
                r = m;
            }
        }

        return l;
    }
}

class BinaryIndexedTree {
    int n;
    vector<long long> c;

public:
    BinaryIndexedTree(int n)
        : n(n)
        , c(n + 1) {}

    void update(int x, long long delta) {
        while (x <= n) {
            c[x] += delta;
            x += x & -x;
        }
    }

    long long query(int x) {
        long long s = 0;
        while (x) {
            s += c[x];
            x -= x & -x;
        }
        return s;
    }
};

class Solution {
public:
    long long distantSubarrays(vector<int>& nums, int goal, int k) {
        int n = nums.size();
        vector<long long> s(n + 1);

        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + nums[i];
        }

        vector<long long> st = s;
        sort(st.begin(), st.end());

        long long ans = 1LL * n * (n + 1) / 2;
        BinaryIndexedTree bit(st.size() + 1);

        for (long long v : s) {
            long long a = v - goal - k + 1LL;
            long long b = v - goal + k - 1LL;

            int l = lower_bound(st.begin(), st.end(), a) - st.begin() + 1;
            int r = lower_bound(st.begin(), st.end(), b + 1) - st.begin();

            if (l <= r) {
                ans -= bit.query(r) - bit.query(l - 1);
            }

            int pos = lower_bound(st.begin(), st.end(), v) - st.begin() + 1;
            bit.update(pos, 1);
        }

        return ans;
    }
};

class BinaryIndexedTree:
    __slots__ = "n", "c"

    def __init__(self, n: int):
        self.n = n
        self.c = [0] * (n + 1)

    def update(self, x: int, delta: int) -> None:
        while x <= self.n:
            self.c[x] += delta
            x += x & -x

    def query(self, x: int) -> int:
        s = 0
        while x:
            s += self.c[x]
            x -= x & -x
        return s


class Solution:
    def distantSubarrays(self, nums: list[int], goal: int, k: int) -> int:
        s = list(accumulate(nums, initial=0))
        st = sorted(s)
        n = len(nums)
        ans = (1 + n) * n // 2
        bit = BinaryIndexedTree(len(st) + 1)
        for v in s:
            a = v - goal - k + 1
            b = v - goal + k - 1

            l = bisect_left(st, a) + 1
            r = bisect_left(st, b + 1)
            if l <= r:
                ans -= bit.query(r) - bit.query(l - 1)
            bit.update(bisect_left(st, v) + 1, 1)
        return ans

type BinaryIndexedTree struct {
	n int
	c []int
}

func NewBinaryIndexedTree(n int) *BinaryIndexedTree {
	return &BinaryIndexedTree{
		n: n,
		c: make([]int, n+1),
	}
}

func (t *BinaryIndexedTree) update(x, delta int) {
	for x <= t.n {
		t.c[x] += delta
		x += x & -x
	}
}

func (t *BinaryIndexedTree) query(x int) int {
	s := 0
	for x > 0 {
		s += t.c[x]
		x -= x & -x
	}
	return s
}

func distantSubarrays(nums []int, goal int, k int) int64 {
	n := len(nums)
	s := make([]int, n+1)

	for i, x := range nums {
		s[i+1] = s[i] + x
	}

	st := append([]int(nil), s...)
	sort.Ints(st)

	ans := n * (n + 1) / 2
	bit := NewBinaryIndexedTree(len(st) + 1)

	for _, v := range s {
		a := v - goal - k + 1
		b := v - goal + k - 1

		l := sort.SearchInts(st, a) + 1
		r := sort.SearchInts(st, b+1)

		if l <= r {
			ans -= bit.query(r) - bit.query(l-1)
		}

		bit.update(sort.SearchInts(st, v)+1, 1)
	}

	return int64(ans)
}

class BinaryIndexedTree {
    private readonly n: number;
    private readonly c: number[];

    constructor(n: number) {
        this.n = n;
        this.c = new Array(n + 1).fill(0);
    }

    update(x: number, delta: number): void {
        while (x <= this.n) {
            this.c[x] += delta;
            x += x & -x;
        }
    }

    query(x: number): number {
        let s = 0;
        while (x > 0) {
            s += this.c[x];
            x -= x & -x;
        }
        return s;
    }
}

function distantSubarrays(nums: number[], goal: number, k: number): number {
    const n = nums.length;
    const s = new Array<number>(n + 1).fill(0);

    for (let i = 0; i < n; i++) {
        s[i + 1] = s[i] + nums[i];
    }

    const st = [...s].sort((a, b) => a - b);

    let ans = (n * (n + 1)) / 2;
    const bit = new BinaryIndexedTree(st.length + 1);

    for (const v of s) {
        const a = v - goal - k + 1;
        const b = v - goal + k - 1;

        const l = _.sortedIndex(st, a) + 1;
        const r = _.sortedIndex(st, b + 1);

        if (l <= r) {
            ans -= bit.query(r) - bit.query(l - 1);
        }

        bit.update(_.sortedIndex(st, v) + 1, 1);
    }

    return ans;
}

4051 - Count Subarrays with Distant Sums

4051. Count Subarrays with Distant Sums

Description

Solutions

Solution 1: Prefix Sum + Binary Indexed Tree

All Problems

All Solutions

Welcome to Subscribe On Youtube

4051. Count Subarrays with Distant Sums

Description

Solutions

Solution 1: Prefix Sum + Binary Indexed Tree

All Problems

All Solutions