extra/rosetta-code/equilibrium-index/equilibrium-index.factor

   1 ! Copyright (c) 2012 Anonymous
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: kernel math sequences ;
   4 IN: rosetta-code.equilibrium-index
   5
   6 ! http://rosettacode.org/wiki/Equilibrium_index
   7
   8 ! An equilibrium index of a sequence is an index into the sequence such that the sum of elements at lower indices is equal to the sum of elements at higher indices. For example, in a sequence A:
   9 !   A0 = − 7
  10 !   A1 = 1
  11 !   A2 = 5
  12 !   A3 = 2
  13 !   A4 = − 4
  14 !   A5 = 3
  15 !   A6 = 0
  16
  17 ! 3 is an equilibrium index, because:
  18 !   A0 + A1 + A2 = A4 + A5 + A6
  19
  20 ! 6 is also an equilibrium index, because:
  21 !   A0 + A1 + A2 + A3 + A4 + A5 = 0
  22 !   (sum of zero elements is zero)
  23
  24 ! 7 is not an equilibrium index, because it is not a valid index
  25 ! of sequence A.
  26
  27 ! Write a function that, given a sequence, returns its
  28 ! equilibrium indices (if any). Assume that the sequence may be
  29 ! very long.
  30
  31 : accum-left ( seq id quot -- seq )
  32     accumulate nip ; inline
  33
  34 : accum-right ( seq id quot -- seq )
  35     [ <reversed> ] 2dip accum-left <reversed> ; inline
  36
  37 : equilibrium-indices ( seq -- inds )
  38     0 [ + ] [ accum-left ] [ accum-right ] 3bi [ = ] 2map
  39     V{ } swap dup length iota [ [ suffix ] curry [ ] if ] 2each ;