I am trying to iterate over an array and count the number of positive, negative and zeros in an array. Right now I am doing it like this
arr = [1, -1, 0, 2, 3, -2, -5]
pos = arr.select { |i| i > 0 }.count
neg = arr.select { |i| i < 0 }.count
zero = arr.select { |i| i == 0 }.count
puts pos
puts neg
puts zero
But is there any way where I can do this in one line? Something like this?
pos, neg, zero = arr.select { |i| i > 0; i < 0; i == 0; }.count
Use inject and the <=> operator:
neg, zero, pos = arr.inject([0,0,0]) { |a,b| a[(b<=>0)+1] += 1; a }
Alternatively, as @HolgerJust mentioned:
neg, zero, pos = arr.each_with_object([0,0,0]) { |a,b| b[(a<=>0)+1] += 1 }
is slightly longer but doesn't have the extra ; a in the block.
Inspired by @steenslag's use of tally:
neg, zero, pos = arr.map { |x| x<=>0 }.tally.values_at(-1,0,1)
each_with_object to avoid the trailing "return value": neg, zero, pos = arr.each_with_object([0,0,0]) { |resullt, i| result[(i<=>0)+1] += 1 }If you use a counting hash the code is short and the results are returned in a hash, which may be convenient.
arr = [1, -1, 0, 2, 3, -2, -5, 4]
You could write
arr.each_with_object(Hash.new(0)) { |n,h| h[n<=>0] += 1 }
#=> {1=>4, -1=>3, 0=>1}
or perhaps you would prefer
labels = { -1=>:neg, 0=>:zero, 1=>:pos }
arr.each_with_object(Hash.new(0)) { |n,h| h[labels[n<=>0]] += 1 }
#=> {:pos=>4, :neg=>3, :zero=>1}
the last line of which could alternatively be written
arr.each_with_object({}) { |n,h| h[labels[n<=>0]] = (h[labels[n<=>0]] ||= 0) + 1 }
See Hash::new, specifically the (second) form that takes an argument called the default value (here zero), and no block. If a hash is defined h = Hash.new(0), then if h has no key k, h[k] returns 0 (and h is not changed).
As others have already said, you should just use inject and count using the <=> operator. If you plan to use similar logic frequently, you could monkey patch a #tally_by method into Enumerable like so:
class Enumerable
def my_tally(*keys, &proc)
proc ||= -> e {e} # Default identity proc
h = keys.empty? ? Hash.new(0) : Hash[keys.map{|k|[k, 0]}]
inject(h){|a, e| a[proc.call(e)] += 1; a}
end
end
This allows you to write:
neg, zero, pos = arr.my_tally(-1, 0, 1){|e| e <=> 0}
While this is certainly more upfront code than the others, it may be nice to have if you find yourself using similar logic frequently. You could also just make this a regular method somewhere if you don't like monkey-patching.
