Hello, codeforces!

In this article, I would like to show you how to reduce time spent on working with for cycle. There're a few ways to do it, I'll show macro-based and expressible in pure c++.

Introduction

I love for cycle. It looks quite simple:

for (init statement; condition; increment)

But you can adapt it to do many great things. To illustrate the power of for, let's look at the simple binary search problem: we have some values l, r and a predicate p which is true on segment [l, x] and false on (x, r), and we are asked to find x. Using for we can implement the solution in basically just one line:

for (int m = (l + r) / 2; abs(r - l) > 1; (p(m) ? l : r) = m, m = (l + r) / 2);

Pretty neat, huh?

After the cycle finishes, you'll get l = x.

Another important note, the (p(m) ? l : r) = m part works because ternary operator returns reference type. It means that (is_left ? left : right) = new_value; assigns new_value either to left or to right, depending on the value of is_left.

You've probably noticed that here I used abs(r - l) > 1 instead of simple l < r. Well, it happened for a reason: that way we can use this implementation even if l > r and the predicate it true on segment [x, l] and false on (r, x).

To make it yet more elegant (and probably less readable), you could notice that operator , evaluates operands and returns the last one, which allows us to transform two statements m = (l + r) / 2 into a single one:

for (int m; m = (l + r) / 2, abs(r - l) > 1; (p(m) ? l : r) = m);

Now, to make it work for various types of p, l, r let's create a template function:

template <class T, class P>
T bin_search(T left, T right, P pred) {
  for (T mid; mid = (left + right) / 2, mid != left && mid != right; (pred(mid) ? left : right) = mid);
  return left;
}

Changing the condition to mid != left && mid != right makes it almost work for doubles, too: you need to provide a simple class that checks if two doubles are equal with given precision:

struct Float {
  double value;

  bool operator!=(const Float &rhs) const;

  Float operator+(const Float &rhs) const;
  Float operator/(const Float &rhs) const;
};

FOR macro

Okay, in the previous section we've seen that you can do some amazing things with for. However, most of the time you'll probably also need to do something very simple, like iterating over [0, n). I've noticed that a lot of programmers use the short form of for cycle for this. Typically it looks like:

#define FOR(i, n) for (int i = 0; i < n; ++i)

Much simpler and shorter than writing raw for! Then you might also wonder how to do a similar shortcut to iterate over [a, b) and find out that you can't override macro by number of arguments. However, we can use 0 instead of O:

#define FOR(i, n)    for (int i = 0; i < n; ++i)
#define F0R(i, a, b) for (int i = a; i < b; ++i)

But what if a > b? Ok, then we need to add additional checks:

#define FOR(i, n)    for (int i = 0; i < n; ++i)
#define F0R(i, a, b) for (int i = a; i != b; (a < b) ? ++i : --i)

There's a potential bug: if you have huge a and b (not fitting into int for example), then you will get integer overflow for i. To resolve it, you could use decltype expression to get the same type as n for FOR. With F0R it's a bit more complicated. First idea is to simply get type from a or b:

#define FOR(i, n)    for (decltype(n) i = 0; i < n; ++i)
#define F0R(i, a, b) for (decltype(a) i = a; i != b; (a < b) ? ++i : --i)

But what will happen if you use F0R(i, 0, n) or F0R(i, n, 0)? The type of 0 is int, and if n is int64_t you'll probably be unhappy with the results. But there's a solution I found while exploring gcd from the standard library:

#define F0R(i, a, b)  for (common_type_t<decltype(a), decltype(b)> i = a; i != b; (a < b) ? ++i : --i)

Okay, we managed to implement both macros, but now have a drawback of needing to write 0 in some cases. I'm not that smart — can remember only one for macro, having more than one can drive me crazy in mistyping. Can we do better? I googled this a bit and found a trick, that helps actually writing FOR in both cases:

#define GET_MACRO_3(_1, _2, _3, NAME, ...) NAME

#define FORN  ...
#define FORAB ...

#define FOR(...) GET_MACRO_3(__VA_ARGS__, FORAB, FORN)(__VA_ARGS__)

How does it work? GET_MACRO_3 simply expands to its 4th argument. If we pass three parameters to FOR, then the 4th parameter of GET_MACRO_3 will be FORAB, if we pass two, then FORN, thus parameters will be passed to the corresponding macro.

Let's get it all together:

#define GET_MACRO_3(_1, _2, _3, NAME, ...) NAME

#define FORN(i, n)     for (decltype(n) i = 0; i < n; ++i)
#define FORAB(i, a, b) for (common_type_t<decltype(a), decltype(b)> i = a; i != b; (a < b) ? ++i : --i)

#define FOR(...) GET_MACRO_3(__VA_ARGS__, FORAB, FORN)(__VA_ARGS__)

Alternatives to macro

Yeah, having FOR(i, n) is much shorter than for (int i = 0; i < n; ++i), but is there macro-free option? It's the beginning of the section, not the end, so you can guess the answer is yes.

We already have a convenient range-based for. Can we apply it here? There's the straightforward way:

auto range(int n) {
  vector<int> is(n);
  iota(is.begin(), is.end(), 0);
  return is;
}

...

for (int i : range(n))

But there's an issue: this function allocates extra memory to simply iterate over numbers. It seems like mining raisin from pies: you put raisin to pies to extract them only because you have a convenient way to transfer pies. But can we share raisin itself? Ok, let's try to manage it:

for (int i : range(n))

// is equivalent to

auto &&r = range(n);
for (auto begin = r.begin(), end = r.end(); begin != end; ++begin) {
  int i = *begin;
  // loop body
}

So, all we need is to define some object r has methods .begin(), .end() return iterator can ++, *, !=. Ok, let's do it:

struct range_it {
  int value;
  
  int operator*() const { return value; }
  auto &operator++() { ++value; return *this; }
  bool operator!=(const range_it &rhs) const 
  { return value != rhs.value; }
};

struct range {
  int n;

  range() = default;
  explicit range(int n) : n(n) {}

  range_it begin() const { return range_it(0); }
  range_it end() const { return range_it(n); }
};

Ok, here we replaced range(n) from function to class ctor. But what with FORAB? Ok, let's modify it a bit by adding an extra field to range class:

struct range {
  int a, b;

  range() = default;
  explicit range(int n) : a(0), b(n) {}
  explicit range(int a, int b) : a(a), b(b) {}

  range_it begin() const { return range_it(a); }
  range_it end() const { return range_it(b); }
};

Looks good, doesn't it? Here's some working code:

for (int i : range(n))
  for (int j : range(i, n))
    dance_like_no_one_watching(i, j);

But here's a problem: it doesn't work with for (int i : range(n, 0)). Ok, let's check it whenever we need:

struct range_it {
  int value, step;
  
  int operator*() const { return value; }
  auto &operator++() { value += step; return *this; }
  bool operator!=(const range_it &rhs) const 
  { return value != rhs.value; }
};

struct range {
  int a, b, step = 1;

  range() = default;
  range(int n) : a(0), b(n), step(1) {}
  range(int a, int b) : a(a), b(b), step(1) {
    if (a > b) {
      step *= -1;
      swap(a, b);
    }
  }

  range(int a, int b, int step) : a(a), b(b), step(step) {}

  range_it begin() const { return range_it{a, step}; }
  range_it end() const { return range_it{b, step}; }
};

Here we can easily use range as we used to in Python. But what about overflowing? Well, we could've added a template here:

template <class T> struct range_it {
  T value, step;
  
  T operator*() const { return value; }
  auto &operator++() { value += step; return *this; }
  bool operator!=(const range_it &rhs) const 
  { return value != rhs.value; }
};

template <class T> struct range {
  T a, b, step = 1;

  range() = default;
  range(T n) : a(0), b(n), step(1) {}
  range(T a, T b) : a(a), b(b), step(1) {
    if (a > b) {
      step *= -1;
      swap(a, b);
    }
  }

  range(T a, T b, T step) : a(a), b(b), step(step) {}

  range_it begin() const { return range_it{a, step}; }
  range_it end() const { return range_it{b, step}; }
};

But you still need a lot of code to simply iterate through [0, ..., n) backwards: for (int i : range(n-1, -1. -1)). That's a pretty annoying thing to deal with, isn't it? Well, you could've patched range to do it:

template <class T> struct range_it {
  T value, step;
  
  T operator*() const { return value; }
  auto &operator++() { value += step; return *this; }
  bool operator!=(const range_it &rhs) const 
  { return value != rhs.value; }
};

template <class T> struct range {
  T a, b, step = 1;

  range() = default;
  range(T n) : a(0), b(n), step(1) {}
  range(T a, int b) : a(a), b(b), step(1) {
    if (a > b) {
      step *= -1;
      swap(--a, --b);
    }
  }

  range(T a, T b, T step) : a(a), b(b), step(step) {}

  range_it begin() const { return range_it{a, step}; }
  range_it end() const { return range_it{b, step}; }
};

What advantages over a macro?

Since we have the proper range object, we might add any adapters we might need. For example, filter, or transform. Take a look:

bool is_prime(int p) {
  for (int d = 2; d * d <= p; ++d)
    if (p % d == 0) 
      return false;
  return true;
}

int main() {
  for (int i : range(100) 
             | filter(is_prime) 
             | transform([](int x) { return x * x; }))
    cout << i << ' ';
}

That might take some time to implement, you could check one of the possible implementations here. However, if you don't want to do it yourself, it's done in ranges library coming with c++20.