CSES - Subarray Sums II

Authors: Qi Wang, Brad Ma

Edit This Page

Problem

We are asked to find the number of subarrays that sum up to xx 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 ii, we can count the number of prefixes with sum equal to prefixSum[i]−x\texttt{prefixSum}[i] - x. This will ensure that we can remove a prefix from our current prefix to build a subarray with sum xx. After every iteration, we just add our new prefix sum to the map.

Implementation

C++

Time Complexity: O(Nlog⁡N)\mathcal{O}(N\log N)

In C++, because std::unordered_map is vulnerable to collisions, we use std::map at the cost of adding an extra log factor to the complexity.

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


using namespace std;


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

Java

Time Complexity: O(Nlog⁡N)\mathcal{O}(N\log N)

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

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

Warning!

Without incorporating randomness, the code below will fail to pass cases specifically designed to make Python dict run slowly. See the Introduction to Sets module for more information.

import random


RANDOM = random.randrange(2**62)




def Wrapper(x):
	return x ^ RANDOM




def 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