Short answer:
Oxygen atoms (not oxygen molecules) cause the dominant red and green auroral emissions at high altitudes. Although there is more atomic $\ce{O}$ at higher altitudes than nitrogen, the key to understand the different colours is the excitation energy. Auroral particles with higher energies penetrate deeper down into the atmosphere causing higher energetic emissions at lower altitudes. The excitation leading to the red $\ce{O}$ emission higher up needs less energy than the green $\ce{O}$ emission further down. Even more so, the purple $\ce{N2+}$ emission at the very bottom can only be excited by very energetic particles.
Some more details:
You are right to note that molecular oxygen ($\ce{O2}$) is heavier than molecular nitrogen ($\ce{N2}$). Hence, at altitudes above the turbopause at around 105 km, where turbulent mixing ceases to exist because the atmosphere gets too thin, there is relatively more $\ce{N2}$ than $\ce{O2}$ and the ratio between the two increases with altitude. However, there is also strong photodissociation of molecular $\ce{O2}$ to atomic $\ce{O}$. As the recombination of two $\ce{O}$ atoms is slow at high altitudes and the atoms are light, their concentration increases with altitude. Above about 200 km atomic $\ce{O}$ becomes more abundant than $\ce{N2}$.
As stated before, auroral particles with higher energies deposit more of their energy at lower altitudes where they encounter more collisions with the air. Lower energetic particles don't get that far.
The dominant red and green emissions of the aurora are basic excitions of the $\ce{O}$ atom. The red emission at around 630 nm peaks at above 200 km. altitude. It's caused by the excited state $\ce{O}^1D$ with an energy of $\sim5.6$ eV. The green emission at around 558 nm peaks at around 110 km, considerably lower than the red emission. The corresponding state $\ce{O}^1S$ has a higher excitation energy of around $10$ eV. Peaking even lower at around 90 km altitude, the violet $\ce{N2+} $ emission has a wavelength around 428 nm and an excitation energy of around $100$ eV.
The combination of the concentrations of the chemical species and, most importantly, the altitude dependent energy deposition result in the altitude dependence of the colours of the aurora. The further to the violet the light, the higher its energy per photon.
Note that these discussions about what colour is caused by what atom or molecule at what altitude are somewhat simplified. This all only refers to the dominant emission at a certain peak altitude which may overpower fainter colours. To a certain degree, most colours will be emitted from most altitudes, but not be visible to the ground observer. And there are many more colours emitted than the ones discussed.