"""Equilibrium index""" from itertools import (accumulate) # equilibriumIndices :: [Num] -> [Int] def equilibriumIndices(xs): '''List indices at which the sum of values to the left equals the sum of values to the right.''' def go(xs): '''Left scan from accumulate, right scan derived from left''' ls = list(accumulate(xs)) n = ls[-1] return [i for (i, (x, y)) in enumerate(zip( ls, [n] + [n - x for x in ls[0:-1]] )) if x == y] return go(xs) if xs else [] # TEST ------------------------------------------------- # main :: IO () def main(): '''Tabulated test results''' print( tabulated('Equilibrium indices:\n')( equilibriumIndices )([ [-7, 1, 5, 2, -4, 3, 0], [2, 4, 6], [2, 9, 2], [1, -1, 1, -1, 1, -1, 1], [1], [] ]) ) # GENERIC ------------------------------------------------- # tabulated :: String -> (a -> b) -> [a] -> String def tabulated(s): '''heading -> function -> input List -> tabulated output string''' def go(f, xs): def width(x): return len(str(x)) w = width(max(xs, key=width)) return s + '\n' + '\n'.join([ str(x).rjust(w, ' ') + ' -> ' + str(f(x)) for x in xs ]) return lambda f: lambda xs: go(f, xs) if __name__ == '__main__': main()