CSES - Subarray Sums II

Authors: Qi Wang, Brad Ma

Edit This Page

Appears In

Silver - Introduction to Prefix Sums

View Problem Statement

Problem Explanation Implementation

Problem

We are asked to find the number of subarrays that sum up to $x$ given the size of the array and its elements.

Explanation

We can have a map that keeps track of the prefix sums. At each index $i$ , we can count the number of prefixes with sum equal to $\texttt{prefixSum}[i] - x$ . This will ensure that we can remove a prefix from our current prefix to build a subarray with sum $x$ . After every iteration, we just add our new prefix sum to the map.

Implementation

Time Complexity: $\mathcal{O}(N)$

C++

#include <iostream>
#include <map>
#include <vector>

using namespace std;

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);

Java

import java.io.*;
import java.util.*;

public class SubarraySumsII {
	public static void main(String[] args) throws IOException {
		Kattio io = new Kattio();
		int arraySize = io.nextInt();
		int target = io.nextInt();
		int[] array = new int[arraySize];

Python

def main():
	N, X = map(int, input().split())
	prefix, res = 0, 0
	mp = {0: 1}  # mp[0] = 1

	for x in input().split():
		prefix += int(x)
		res += mp.get(prefix - X, 0)  # if not in dict, return 0
		mp[prefix] = mp.get(prefix, 0) + 1
	print(res)


main()

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!

Join Forum