Birthday Chocolate · Coding Gym

Solutions

The solution to this problem when d=1 can be used to solve Picking Numbers as well.

A linear solution consists in:

computing the prefix sum of the input array
keeping a window of size m
counting how many times the window bounds differ by d

Here is a solution in C++:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


int n, d, m;
cin >> n;
vector<int> v(n+1);
copy_n(istream_iterator<int>(cin), n, next(begin(v)));
cin >> d >> m;
partial_sum(begin(v), end(v), begin(v));

auto cnt = 0;
for (auto i=0; i<=n-m; ++i)
{
    cnt += v[i+m] - v[i] == d;
}
cout << cnt;

The loop can be replaced with zip | map | reduce combination of patterns:

1
2
3
4
5
6
7


int n, d, m;
cin >> n;
vector<int> v(n+1);
copy_n(istream_iterator<int>(cin), n, next(begin(v)));
cin >> d >> m;
partial_sum(begin(v), end(v), begin(v));
cout << inner_product(next(begin(v), m), end(v), begin(v), 0, plus<>{}, [=](auto l, auto r){ return l-r == d; });

Here is an alternative solution in Python filtering out non-matching pairs:

1
2
3
4
5
6


def solve(n, s, d, m):
    for i in range(1, len(s)):
        s[i] += s[i-1]
    s = zip(s, s[m:])
    s = filter(lambda p: p[1]-p[0]==d, s)
    return len(s)

Another interesting solution with C++ ranges:

1
2
3


auto psum = view::partial_sum(v);
auto zipped = view::zip_with([=](auto l, auto r) { return r - l == d; }, psum, view::drop(psum, m));    
cout << ranges::accumulate(zipped, 0);

Note: a brute force solution also works. In particular, please note that m is at most 31 so the internal loop of the following code (courtesy of HackerRank) is time-constant:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29


int n;
cin >> n;
vector<int> c(n);
for(int i = 0; i < n; i++)
{
   cin >> c[i];
}
int d;
int m;
cin >> d >> m;

int ans = 0;
for(int i = 0; i < n; i++)
{
    if (i + m - 1 < n)
    {
        int sum = 0;
        // the following loop has at most m iterations (m is a constant of the problem)
        for (int j = i; j <= i + m - 1; j++)
        {
            sum = sum + c[j];
        }
        if (sum==d) 
        {
            ans++;
        }
    }
}
cout << ans << endl;

An Haskell solution.

1
2
3
4
5
6
7
8


-- Count segments where the sum of elements is d
solve :: ([Int], Int, Int) -> Int
solve (s, d, m) = length $ filter (== d) $ sum <$> segments m s

-- Return all contiguous segments of length n from the given list.
segments :: Int -> [a] -> [[a]]
segments _ [] = []
segments n xs = take n xs : segments n (drop 1 xs)

The segments n xs function generates a list of segments of length n:

segments n [2,2,1,3,2] -> [2,2], [2,1], [1,3], [3,2]

And solve function:

sums elements: [4, 3, 4, 5]
filters only elements equal to d=4: [4, 4]
count them: 2

prefix sum patterns