Equal Stacks · Coding Gym

Solutions

C++ Solution very pretty much independent from the number of the stacks (just change numberOfStacks):

 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


const int numberOfStacks = 3;

array<int, numberOfStacks> n{};
for (auto& i : n)
    cin >> i;

array<vector<int>, numberOfStacks> stacks;
array<int, numberOfStacks> heights{};

for (auto i=0; i<n.size(); ++i)
{
    auto& s = stacks[i];
    s.resize(n[i]);
    copy_n(istream_iterator<int>(cin), n[i], begin(s));
    reverse(begin(s), end(s));
    // we may even accumulate while we read 
    heights[i] = accumulate(begin(s), end(s), 0);
}

while (!equal(next(begin(heights)), end(heights), begin(heights)))
{
    auto maxIt = max_element(begin(heights), end(heights));
    auto maxIdx = distance(begin(heights), maxIt);
    heights[maxIdx] -= stacks[maxIdx].back();
    stacks[maxIdx].pop_back();
}

cout << heights.front();

Basically, we first calculate the heights of all the stacks (sum of the elements of each) and then we remove one element from the tallest until all the heights are the same.

A similar solution by Alessandro Pezzato, implemented in Haskell with recursion:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


import Data.List (sortOn)

allTheSame :: [Int] -> Bool
allTheSame = and . (zipWith (==) <*> tail)

cutHigher :: [[Int]] -> [[Int]]
cutHigher stacks = tail higher : other
  where
    (higher:other) = sortOn (((-1) *) . sum) stacks

equalStacks :: [[Int]] -> Int
equalStacks stacks =
  if allTheSame (fmap sum stacks)
    then sum (head stacks)
    else equalStacks $ cutHigher stacks

main =
  interact (show . equalStacks . (fmap . fmap) read . fmap words . tail . lines)

But the solution above has a problem: the stack height is calculated every time. We can optimize: calculate the height only when constructing the stack, then subtracting at each pop.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


data Stack = Stack (Int, [Int])

fromString str = Stack (sum elements, elements)
  where
    elements = fmap read $ words str

height (Stack (h, _)) = h

allTheSame = and . (zipWith (==) <*> tail)

pop (Stack (h, (x:xs))) = Stack (h - x, xs)

cutHigher stacks = pop higher : other
  where
    (higher:other) = sortOn (negate . height) stacks

equalStacks stacks =
  if allTheSame (fmap height stacks)
    then height $ head stacks
    else equalStacks $ cutHigher stacks

main = interact (show . equalStacks . fmap fromString . tail . lines)

Another similar solution by Francesco Dammacco, implemented in JavaScript:

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


function equalStacks(stacks) {
    stacks.forEach(stack => stack.reverse());
    var heights = stacks.map(stack => stack.reduce((accumulator, element) => accumulator + element));
       
    while(!heights.every( height => height === heights[0] )) {        
        var maxSum = Math.max(...heights);
        var maxIndex = heights.indexOf(maxSum);
        heights[maxIndex] -= stacks[maxIndex].pop();
    }
    
    return heights[0];
}

Prefix sum solutions

An alternative solution is based on the concept of prefix sum. Basically, we can calculate the prefix sum of all the stacks and find the maximum intersection:

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


n1, n2, n3 = map(int, raw_input().split())
s1 = [0] + map(int, raw_input().split())[::-1]
s2 = [0] + map(int, raw_input().split())[::-1]
s3 = [0] + map(int, raw_input().split())[::-1]

for i in xrange(1, len(s1)):
    s1[i] += s1[i-1]
for i in xrange(1, len(s2)):
    s2[i] += s2[i-1]
for i in xrange(1, len(s3)):
    s3[i] += s3[i-1]
print max(set(s1) & set(s2) & set(s3))

This is an adaptation of the haskell solution above, using prefix sum. The stack front is the current height, so we can recursively remove the higher value until all the front are the same. We must add a 0 at the beginning, so we don’t end up with empty lists, causing an exception on head.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


import Data.List (sortOn)

fromString = reverse . (0:) . prefixSum . reverse . fmap read . words

allTheSame = and . (zipWith (==) <*> tail)

cutHigher stacks = tail higher : other
  where
    higher : other = sortOn (negate . head) stacks

prefixSum = scanl1 (+)

equalStacks stacks =
  if allTheSame $ fmap head stacks
    then head . head $ stacks
    else equalStacks $ cutHigher stacks

main = interact $ show . equalStacks . fmap fromString . tail . lines

prefix sum