2007/12/3, Yusuke ENDOH <mame / tsg.ne.jp>:
> Hi,
>
> 2007/12/3, Voroztsov Artem <artem.voroztsov / gmail.com>:
> > I tried to add missed functionality:
>
> It seems to not be able to handle a non-recursive array that
> contains multiple references to one another array:
>
>
>   a = []
>   a = [a, a]
>   p a.flatten2  #=> tried to flatten recursive array
>
>


OK, i see.
I fixed bug in my ruby version  (it still not recursive :))) ) (just
for fun and for quick patching).

It is 4 times slower than original flatten for this sample:

a = [0]; 20.times { a = [a, a] }; a.flatten

but still much faster for long depth nested arrays.

And I wrote some unit tests.
Core team may want to use them.


#!/usr/bin/ruby

  class Array
    def flatten2(level = nil)
      return flatten3(level) if level
      res = []
      stack = [[self.dup, nil]]
      met = {}
      while !stack.empty?
        a,top_ary = stack.pop
        if top_ary
          met.delete(top_ary.object_id)
          next
        end
        if a.is_a?(Array)
          raise ArgumentError, "tried to flatten recursive array" if
met[a.object_id]
          unless a.empty?
            met[a.object_id] = 1
            stack << [nil, a] # forget_me statement
            a.each{|i| stack << [i]}
          end
        else
          res << a
        end
      end
      res.reverse
    end

    private
    def flatten3(level)
      res = []
      stack = [[self.dup, level+1, nil]] # stack to emulatine recursive calls
      met = {}
      while !stack.empty?
        a,level,top_ary = stack.pop
        if top_ary
          met.delete(top_ary.object_id)
          next
        end
        if a.is_a?(Array) && level > 0
          raise ArgumentError, "tried to flatten recursive array" if
met[a.object_id]
          unless a.empty?
            stack << [nil, nil, a] # forget_me statement
            met[a.object_id] = 1
            a.each{|i| stack << [i,level-1]}
          end
        else
          res << a
        end
      end
      res.reverse
    end
  end


  if $0 == __FILE__
    require 'test/unit'
    require 'benchmark'

    a = [0]; 15.times { a = [a, a] }; a.flatten2

    Benchmark.bm {|b|
      b.report('original flatten') {
        res1 = a.flatten
      }
      b.report('self-made flatten') {
        res2 = a.flatten2
      }
    }

    $method = :flatten2

    class TestFlatten < Test::Unit::TestCase
      @@samples = [
          [ [],[] ],
          [ [:abc],[:abc] ],
          [ [[[1]]], [1]  ],
          [ [1,2,3], [1,2,3] ],
          [ [[[[1],2],3],4], [1,2,3,4] ],
          [ [1,[2,[3,[4,[5]]]]], [1,2,3,4,5] ],
          [ [:a,[:b,[:c],:d],:e], [:a,:b,:c,:d,:e]] ,
          [ [[1,[2,3],4,[5,[6,[7]],[8,9]],10]], (1..10).to_a ],
          [ [[],[[]],[]], [] ],
          [ ['a',[],[['b']],[], 'c'], ['a', 'b', 'c']]
      ]

      def stest_simple
        @@samples.each do |ary ,ans|
          assert_equal(ans, ary.send($method))
        end
      end

      def test_with_arg
        x = [1,[2,[3,[4,5]]]]
        ans = [
          [1,[2,[3,[4,5]]]],
          [1,2,[3,[4,5]]],
          [1,2,3,[4,5]],
          [1,2,3,4,5],
          [1,2,3,4,5]
        ]
        5.times do |i|
          assert_equal(ans[i], x.send($method,i))
        end

        x = [[[[1,2],3],4],5]
        ans = [
          [[[[1,2],3],4],5],
          [[[1,2],3],4,5],
          [[1,2],3,4,5],
          [1,2,3,4,5],
          [1,2,3,4,5]
        ]
        5.times do |i|
          assert_equal(ans[i], x.send($method,i))
        end
      end

      def test_tree
        n = 10
        a = [0];
        n.times { a = [a, a] };
        assert_equal(Array.new(2**n){0}, a.send($method))
      end

      def test_depth
        n = 10000
        a = []; n.times {|i| a = [a, i] };
        ans = (0...n).to_a
        assert_equal( ans, a.send($method))

        a = []; n.times {|i| a = [i,a] };
        assert_equal( ans.reverse, a.send($method))

        a = []; n.times {|i| a = [i, a, i] };

        assert_equal( ans.reverse + ans, a.send($method))
      end

      def test_recursive
        i = 0
        a = [[1]]
        a = [a,a]
        assert_equal([1,1], a.send($method))

        a << a
        assert_raise ArgumentError do
          a.send($method)
        end

        a = [1,[2,3,[4,5]]]
        a.last.last << a
        assert_raise ArgumentError do
          a.send($method)
        end

        a = [1,2]
        b = [3,4]
        a << b
        b << [5,6]
        c = [7,a,8,b,9]
        assert_equal([7,1,2,3,4,5,6,8,3,4,5,6,9], c.send($method))
        c1 = [a,b]
        b << c1
        assert_raise ArgumentError do
          c1.send($method)
        end
        b << a
        b << 10
        assert_raise ArgumentError do
          c.send($method)
        end
      end
    end
  end

############################
OUTPUT

>ruby test_flatten3.rb
      user     system      total        real
original flatten  6.391000   0.297000   6.688000 (  6.703000)
self-made flatten 23.640000   0.000000  23.640000 ( 23.641000)
Loaded suite test_flatten3
Started
....
Finished in 0.39 seconds.

4 tests, 20 assertions, 0 failures, 0 errors